专栏首页Dechin的专栏Python3实现打格点算法的GPU加速

Python3实现打格点算法的GPU加速

技术背景

在数学和物理学领域,总是充满了各种连续的函数模型。而当我们用现代计算机的技术去处理这些问题的时候,事实上是无法直接处理连续模型的,绝大多数的情况下都要转化成一个离散的模型再进行数值的计算。比如计算数值的积分,计算数值的二阶导数(海森矩阵)等等。这里我们所介绍的打格点的算法,正是一种典型的离散化方法。这个对空间做离散化的方法,可以在很大程度上简化运算量。比如在分子动力学模拟中,计算近邻表的时候,如果不采用打格点的方法,那么就要针对整个空间所有的原子进行搜索,计算出来距离再判断是否近邻。而如果采用打格点的方法,我们只需要先遍历一遍原子对齐进行打格点的离散化,之后再计算近邻表的时候,只需要计算三维空间下邻近的27个格子中的原子是否满足近邻条件即可。在这篇文章中,我们主要探讨如何用GPU来实现打格点的算法。

打格点算法实现

我们先来用一个例子说明一下什么叫打格点。对于一个给定所有原子坐标的系统,也就是已知了

[x,y,z]

,我们需要得到的是这些原子所在的对应的格子位置

[n_x,n_y,n_z]

。我们先看一下在CPU上的实现方案,是一个遍历一次的算法:

# cuda_grid.py

from numba import jit
from numba import cuda
import numpy as np

def grid_by_cpu(crd, rxyz, atoms, grids):
    """Transform coordinates [x,y,z] into grids [nx,ny,nz].
    Args:
        crd(list): The 3-D coordinates of atoms.
        rxyz(list): The list includes xmin,ymin,zmin,grid_num.
        atoms(int): The total number of atoms.
        grids(list): The transformed grids matrix.
    """
    for i in range(atoms):
        grids[i][0] = int((crd[i][0]-rxyz[0])/rxyz[3])
        grids[i][1] = int((crd[i][1]-rxyz[1])/rxyz[3])
        grids[i][2] = int((crd[i][2]-rxyz[2])/rxyz[3])
    return grids

if __name__=='__main__':
    np.random.seed(1)
    atoms = 4
    grid_size = 0.1
    crd = np.random.random((atoms,3)).astype(np.float32)
    xmin = min(crd[:,0])
    ymin = min(crd[:,1])
    zmin = min(crd[:,2])
    xmax = max(crd[:,0])
    ymax = max(crd[:,1])
    zmax = max(crd[:,2])
    xgrids = int((xmax-xmin)/grid_size)+1
    ygrids = int((ymax-ymin)/grid_size)+1
    zgrids = int((zmax-zmin)/grid_size)+1
    rxyz = np.array([xmin,ymin,zmin,grid_size], dtype=np.float32)
    
    grids = np.ones_like(crd)*(-1)
    grids = grids.astype(np.float32)
    grids_cpu = grid_by_cpu(crd, rxyz, atoms, grids)
    print (crd)
    print (grids_cpu)

    import matplotlib.pyplot as plt
    plt.figure()
    plt.plot(crd[:,0], crd[:,1], 'o', color='red')
    for grid in range(ygrids+1):
        plt.plot([xmin,xmin+grid_size*xgrids], [ymin+grid_size*grid,ymin+grid_size*grid], color='black')
    for grid in range(xgrids+1):
        plt.plot([xmin+grid_size*grid,xmin+grid_size*grid], [ymin,ymin+grid_size*ygrids], color='black')
    plt.savefig('Atom_Grids.png')

输出结果如下,

$ python3 cuda_grid.py 
[[4.17021990e-01 7.20324516e-01 1.14374816e-04]
 [3.02332580e-01 1.46755889e-01 9.23385918e-02]
 [1.86260208e-01 3.45560730e-01 3.96767467e-01]
 [5.38816750e-01 4.19194520e-01 6.85219526e-01]]
[[2. 5. 0.]
 [1. 0. 0.]
 [0. 1. 3.]
 [3. 2. 6.]]

上面两个打印输出就分别对应于

[x,y,z]

[n_x,n_y,n_z]

,比如第一个原子被放到了编号为

[2,5,0]

的格点。那么为了方便理解打格点的方法,我们把这个三维空间的原子系统和打格点以后的标号取前两个维度来可视化一下结果,作图以后效果如下:

我们可以看到,这些红色的点就是原子所处的位置,而黑色的网格线就是我们所标记的格点。在原子数量比较多的时候,有可能出现在一个网格中存在很多个原子的情况,所以如何打格点,格点大小如何去定义,这都是不同场景下的经验参数,需要大家一起去摸索。

打格点算法加速

在上面这个算法实现中,我们主要是用到了一个for循环,这时候我们可以想到numba所支持的向量化运算,还有GPU硬件加速,这里我们先对比一下三种实现方案的计算结果:

# cuda_grid.py

from numba import jit
from numba import cuda
import numpy as np

def grid_by_cpu(crd, rxyz, atoms, grids):
    """Transform coordinates [x,y,z] into grids [nx,ny,nz].
    Args:
        crd(list): The 3-D coordinates of atoms.
        rxyz(list): The list includes xmin,ymin,zmin,grid_num.
        atoms(int): The total number of atoms.
        grids(list): The transformed grids matrix.
    """
    for i in range(atoms):
        grids[i][0] = int((crd[i][0]-rxyz[0])/rxyz[3])
        grids[i][1] = int((crd[i][1]-rxyz[1])/rxyz[3])
        grids[i][2] = int((crd[i][2]-rxyz[2])/rxyz[3])
    return grids

@jit
def grid_by_jit(crd, rxyz, atoms, grids):
    """Transform coordinates [x,y,z] into grids [nx,ny,nz].
    Args:
        crd(list): The 3-D coordinates of atoms.
        rxyz(list): The list includes xmin,ymin,zmin,grid_num.
        atoms(int): The total number of atoms.
        grids(list): The transformed grids matrix.
    """
    for i in range(atoms):
        grids[i][0] = int((crd[i][0]-rxyz[0])/rxyz[3])
        grids[i][1] = int((crd[i][1]-rxyz[1])/rxyz[3])
        grids[i][2] = int((crd[i][2]-rxyz[2])/rxyz[3])
    return grids

@cuda.jit
def grid_by_gpu(crd, rxyz, grids):
    """Transform coordinates [x,y,z] into grids [nx,ny,nz].
    Args:
        crd(list): The 3-D coordinates of atoms.
        rxyz(list): The list includes xmin,ymin,zmin,grid_num.
        atoms(int): The total number of atoms.
        grids(list): The transformed grids matrix.
    """
    i,j = cuda.grid(2)
    grids[i][j] = int((crd[i][j]-rxyz[j])/rxyz[3])

if __name__=='__main__':
    np.random.seed(1)
    atoms = 4
    grid_size = 0.1
    crd = np.random.random((atoms,3)).astype(np.float32)
    xmin = min(crd[:,0])
    ymin = min(crd[:,1])
    zmin = min(crd[:,2])
    xmax = max(crd[:,0])
    ymax = max(crd[:,1])
    zmax = max(crd[:,2])
    xgrids = int((xmax-xmin)/grid_size)+1
    ygrids = int((ymax-ymin)/grid_size)+1
    zgrids = int((zmax-zmin)/grid_size)+1
    rxyz = np.array([xmin,ymin,zmin,grid_size], dtype=np.float32)
    crd_cuda = cuda.to_device(crd)
    rxyz_cuda = cuda.to_device(rxyz)
    
    grids = np.ones_like(crd)*(-1)
    grids = grids.astype(np.float32)
    grids_cpu = grid_by_cpu(crd, rxyz, atoms, grids)

    grids = np.ones_like(crd)*(-1)
    grids_jit = grid_by_jit(crd, rxyz, atoms, grids)

    grids = np.ones_like(crd)*(-1)
    grids_cuda = cuda.to_device(grids)
    
    grid_by_gpu[(atoms,3),(1,1)](crd_cuda,
                                 rxyz_cuda,
                                 grids_cuda)

    print (crd)
    print (grids_cpu)
    print (grids_jit)
    print (grids_cuda.copy_to_host())

输出结果如下:

$ python3 cuda_grid.py 
/home/dechin/anaconda3/lib/python3.8/site-packages/numba/cuda/compiler.py:865: NumbaPerformanceWarning: Grid size (12) < 2 * SM count (72) will likely result in GPU under utilization due to low occupancy.
  warn(NumbaPerformanceWarning(msg))
[[4.17021990e-01 7.20324516e-01 1.14374816e-04]
 [3.02332580e-01 1.46755889e-01 9.23385918e-02]
 [1.86260208e-01 3.45560730e-01 3.96767467e-01]
 [5.38816750e-01 4.19194520e-01 6.85219526e-01]]
[[2. 5. 0.]
 [1. 0. 0.]
 [0. 1. 3.]
 [3. 2. 6.]]
[[2. 5. 0.]
 [1. 0. 0.]
 [0. 1. 3.]
 [3. 2. 6.]]
[[2. 5. 0.]
 [1. 0. 0.]
 [0. 1. 3.]
 [3. 2. 6.]]

我们先看到这里面的告警信息,因为GPU硬件加速要在一定密度的运算量之上才能够有比较明显的加速效果。比如说我们只是计算两个数字的加和,那么是完全没有必要使用到GPU的。但是如果我们要计算两个非常大的数组的加和,那么这个时候GPU就能够发挥出非常大的价值。因为这里我们的案例中只有4个原子,因此提示我们这时候是体现不出来GPU的加速效果的。我们仅仅关注下这里的运算结果,在不同体系下得到的格点结果是一致的,那么接下来就可以对比一下几种不同实现方式的速度差异。

# cuda_grid.py

from numba import jit
from numba import cuda
import numpy as np

def grid_by_cpu(crd, rxyz, atoms, grids):
    """Transform coordinates [x,y,z] into grids [nx,ny,nz].
    Args:
        crd(list): The 3-D coordinates of atoms.
        rxyz(list): The list includes xmin,ymin,zmin,grid_num.
        atoms(int): The total number of atoms.
        grids(list): The transformed grids matrix.
    """
    for i in range(atoms):
        grids[i][0] = int((crd[i][0]-rxyz[0])/rxyz[3])
        grids[i][1] = int((crd[i][1]-rxyz[1])/rxyz[3])
        grids[i][2] = int((crd[i][2]-rxyz[2])/rxyz[3])
    return grids

@jit
def grid_by_jit(crd, rxyz, atoms, grids):
    """Transform coordinates [x,y,z] into grids [nx,ny,nz].
    Args:
        crd(list): The 3-D coordinates of atoms.
        rxyz(list): The list includes xmin,ymin,zmin,grid_num.
        atoms(int): The total number of atoms.
        grids(list): The transformed grids matrix.
    """
    for i in range(atoms):
        grids[i][0] = int((crd[i][0]-rxyz[0])/rxyz[3])
        grids[i][1] = int((crd[i][1]-rxyz[1])/rxyz[3])
        grids[i][2] = int((crd[i][2]-rxyz[2])/rxyz[3])
    return grids

@cuda.jit
def grid_by_gpu(crd, rxyz, grids):
    """Transform coordinates [x,y,z] into grids [nx,ny,nz].
    Args:
        crd(list): The 3-D coordinates of atoms.
        rxyz(list): The list includes xmin,ymin,zmin,grid_num.
        atoms(int): The total number of atoms.
        grids(list): The transformed grids matrix.
    """
    i,j = cuda.grid(2)
    grids[i][j] = int((crd[i][j]-rxyz[j])/rxyz[3])

if __name__=='__main__':
    import time
    from tqdm import trange

    np.random.seed(1)
    atoms = 100000
    grid_size = 0.1
    crd = np.random.random((atoms,3)).astype(np.float32)
    xmin = min(crd[:,0])
    ymin = min(crd[:,1])
    zmin = min(crd[:,2])
    xmax = max(crd[:,0])
    ymax = max(crd[:,1])
    zmax = max(crd[:,2])
    xgrids = int((xmax-xmin)/grid_size)+1
    ygrids = int((ymax-ymin)/grid_size)+1
    zgrids = int((zmax-zmin)/grid_size)+1
    rxyz = np.array([xmin,ymin,zmin,grid_size], dtype=np.float32)
    crd_cuda = cuda.to_device(crd)
    rxyz_cuda = cuda.to_device(rxyz)
    
    cpu_time = 0
    jit_time = 0
    gpu_time = 0

    for i in trange(100):
        grids = np.ones_like(crd)*(-1)
        grids = grids.astype(np.float32)
        time0 = time.time()
        grids_cpu = grid_by_cpu(crd, rxyz, atoms, grids)
        time1 = time.time()

        grids = np.ones_like(crd)*(-1)
        time2 = time.time()
        grids_jit = grid_by_jit(crd, rxyz, atoms, grids)
        time3 = time.time()

        grids = np.ones_like(crd)*(-1)
        grids_cuda = cuda.to_device(grids)
        time4 = time.time()
        grid_by_gpu[(atoms,3),(1,1)](crd_cuda,
                                    rxyz_cuda,
                                    grids_cuda)
        time5 = time.time()
        
        if i != 0:
            cpu_time += time1 - time0
            jit_time += time3 - time2
            gpu_time += time5 - time4
    
    print ('The time cost of CPU calculation is: {}s'.format(cpu_time))
    print ('The time cost of JIT calculation is: {}s'.format(jit_time))
    print ('The time cost of GPU calculation is: {}s'.format(gpu_time))

输出结果如下:

$ python3 cuda_grid.py 
100%|███████████████████████████| 100/100 [00:23<00:00,  4.18it/s]
The time cost of CPU calculation is: 23.01943016052246s
The time cost of JIT calculation is: 0.04810166358947754s
The time cost of GPU calculation is: 0.01806473731994629s

在100000个原子的体系规模下,普通的for循环实现效率就非常的低下,需要23s,而经过向量化运算的加速之后,直接飞升到了0.048s,而GPU上的加速更是达到了0.018s,相比于没有GPU硬件加速的场景,实现了将近2倍的加速。但是这还远远不是GPU加速的上限,让我们再测试一个更大的案例:

# cuda_grid.py

from numba import jit
from numba import cuda
import numpy as np

def grid_by_cpu(crd, rxyz, atoms, grids):
    """Transform coordinates [x,y,z] into grids [nx,ny,nz].
    Args:
        crd(list): The 3-D coordinates of atoms.
        rxyz(list): The list includes xmin,ymin,zmin,grid_num.
        atoms(int): The total number of atoms.
        grids(list): The transformed grids matrix.
    """
    for i in range(atoms):
        grids[i][0] = int((crd[i][0]-rxyz[0])/rxyz[3])
        grids[i][1] = int((crd[i][1]-rxyz[1])/rxyz[3])
        grids[i][2] = int((crd[i][2]-rxyz[2])/rxyz[3])
    return grids

@jit
def grid_by_jit(crd, rxyz, atoms, grids):
    """Transform coordinates [x,y,z] into grids [nx,ny,nz].
    Args:
        crd(list): The 3-D coordinates of atoms.
        rxyz(list): The list includes xmin,ymin,zmin,grid_num.
        atoms(int): The total number of atoms.
        grids(list): The transformed grids matrix.
    """
    for i in range(atoms):
        grids[i][0] = int((crd[i][0]-rxyz[0])/rxyz[3])
        grids[i][1] = int((crd[i][1]-rxyz[1])/rxyz[3])
        grids[i][2] = int((crd[i][2]-rxyz[2])/rxyz[3])
    return grids

@cuda.jit
def grid_by_gpu(crd, rxyz, grids):
    """Transform coordinates [x,y,z] into grids [nx,ny,nz].
    Args:
        crd(list): The 3-D coordinates of atoms.
        rxyz(list): The list includes xmin,ymin,zmin,grid_num.
        atoms(int): The total number of atoms.
        grids(list): The transformed grids matrix.
    """
    i,j = cuda.grid(2)
    grids[i][j] = int((crd[i][j]-rxyz[j])/rxyz[3])

if __name__=='__main__':
    import time
    from tqdm import trange

    np.random.seed(1)
    atoms = 5000000
    grid_size = 0.1
    crd = np.random.random((atoms,3)).astype(np.float32)
    xmin = min(crd[:,0])
    ymin = min(crd[:,1])
    zmin = min(crd[:,2])
    xmax = max(crd[:,0])
    ymax = max(crd[:,1])
    zmax = max(crd[:,2])
    xgrids = int((xmax-xmin)/grid_size)+1
    ygrids = int((ymax-ymin)/grid_size)+1
    zgrids = int((zmax-zmin)/grid_size)+1
    rxyz = np.array([xmin,ymin,zmin,grid_size], dtype=np.float32)
    crd_cuda = cuda.to_device(crd)
    rxyz_cuda = cuda.to_device(rxyz)

    jit_time = 0
    gpu_time = 0

    for i in trange(100):
        grids = np.ones_like(crd)*(-1)
        time2 = time.time()
        grids_jit = grid_by_jit(crd, rxyz, atoms, grids)
        time3 = time.time()

        grids = np.ones_like(crd)*(-1)
        grids_cuda = cuda.to_device(grids)
        time4 = time.time()
        grid_by_gpu[(atoms,3),(1,1)](crd_cuda,
                                     rxyz_cuda,
                                     grids_cuda)
        time5 = time.time()
        
        if i != 0:
            jit_time += time3 - time2
            gpu_time += time5 - time4
    
    print ('The time cost of JIT calculation is: {}s'.format(jit_time))
    print ('The time cost of GPU calculation is: {}s'.format(gpu_time))

在这个5000000个原子的案例中,因为普通的for循环已经实在是跑不动了,因此我们就干脆不统计这一部分的时间,最后输出结果如下:

$ python3 cuda_grid.py 
100%|███████████████████████████| 100/100 [00:09<00:00, 10.15it/s]
The time cost of JIT calculation is: 2.3743042945861816s
The time cost of GPU calculation is: 0.022843599319458008s

在如此大规模的运算下,GPU实现100倍的加速,而此时作为对比的CPU上的实现方法是已经用上了向量化运算的操作,也已经可以认为是一个极致的加速了。

总结概要

在这篇文章中,我们主要介绍了打格点算法在分子动力学模拟中的重要价值,以及几种不同的实现方式。其中最普通的for循环的实现效率比较低下,从算法复杂度上来讲却已经是极致。而基于CPU上的向量化运算的技术,可以对计算过程进行非常深度的优化。当然,这个案例在不同的硬件上也能够发挥出明显不同的加速效果,在GPU的加持之下,可以获得100倍以上的加速效果。这也是一个在Python上实现GPU加速算法的一个典型案例。

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

我来说两句

0 条评论
登录 后参与评论

相关文章

  • Python实现GPU加速的基本操作

    之前写过一篇讲述如何使用pycuda来在Python上写CUDA程序的博客。这个方案的特点在于完全遵循了CUDA程序的写法,只是支持了一些常用函数的接口,如果你...

    DechinPhy
  • 超过Numpy的速度有多难?试试Numba的GPU加速

    Numpy是在Python中非常常用的一个库,不仅具有良好的接口文档和生态,还具备了最顶级的性能,这个库很大程度上的弥补了Python本身性能上的缺陷。虽然我们...

    DechinPhy
  • 初识人工智能

    刚开始接触人工智能的时候,大家肯定看到了一些名词:人工智能、深度学习、机器学习...what??哈哈,先来简单的给大家解释一下这三者的区别,再来谈论其他的问题。...

    小闫同学啊
  • python3使用pycuda执行简单GPU计算任务

    GPU的加速技术在深度学习、量子计算领域都已经被广泛的应用。其适用的计算模型是小内存的密集型计算场景,如果计算的模型内存较大,则需要使用到共享内存,这会直接导致...

    DechinPhy
  • Ubuntu中配置TensorFlow使用环境的方法

    TensorFlow™是一个基于数据流编程(dataflow programming)的符号数学系统,被广泛应用于各类机器学习(machine learning...

    砸漏
  • TensorFlow环境搭建

    前几天,我们刚下发了毕业设计的题目,我选的题目为基于TensorFlow的深度学习与研究,这将会是一个系列文章,截止2020-07我会将所有相关内容用更加通俗易...

    石璞东
  • 第18天:NLP实战(二)——用DNN实现手势识别

      接着上次的项目,主要是为了熟悉我们对NLP知识的实际应用,接着上次对语料的获取以及对数据的简单分析,这次,我们通过对DNN——一个比较简单的深度学习算法来实...

    一计之长
  • 一文上手最新TensorFlow2.0系列(二)

    【磐创AI导读】:本系列文章介绍了与tensorflow的相关知识,包括其介绍、安装及使用等。本篇文章将接着上篇文章继续介绍它的安装及部分使用。查看上篇:文末福...

    磐创AI
  • Ubuntu17.04安装TensorFlow1.2的GPU版本

    摘要:Ubuntu17.04安装TensorFlow1.2的GPU版本。首先验证nvidia显卡,然后安装CUDA Toolkit 8.0,安装cuDNN v5...

    王小雷
  • AlphaGo Zero代码迟迟不开源,TF等不及自己推了一个

    Root 编译自GitHub 量子位 出品 | 公众号 QbitAI TensorFlow官方在GitHub上推了一个AlphaGo Zero的开源代码! 这个...

    量子位
  • 简单几步,轻松完成 GPU 云服务器开发环境搭建

    ? 在深度学习和图形处理等领域,GPU相较于CPU有着数十倍到上百倍的算力,能够为企业提供更高的计算效率及更低廉的IT成本,但同时也有不少研究与开发人员对GP...

    腾讯云计算产品团队
  • 毫秒级检测!你见过带GPU加速的树莓派吗?

    用户1332428
  • 毫秒级检测!你见过带GPU的树莓派吗?

    树莓派3B+英特尔神经计算棒进行高速目标检测 转载请注明作者梦里茶 ? 代码: 训练数据预处理: https://gist.github.com/ahan...

    梦里茶
  • 5分钟配置好你的AI开发环境

    无论是第一次设置TensorFlow的新手数据科学爱好者,还是使用TB级数据的经验丰富的AI工程师,安装库、软件包或者框架总是一个困难又繁琐的过程。但是像Doc...

    磐创AI
  • Google人体图像分割模型Bodypix再次更新,针对Coral开发板优化,720p/30fps流畅运行

    AI开发板果然是各大厂商的必争之地。树莓派4刚发布8GB版,快被遗忘的谷歌Coral开发板就迎来新的机器学习模型,登上了近日GitHub开源热榜。

    量子位
  • 5分钟配置好你的AI开发环境

    无论是第一次设置TensorFlow的新手数据科学爱好者,还是使用TB级数据的经验丰富的AI工程师,安装库、软件包或者框架总是一个困难又繁琐的过程。但是像Doc...

    用户7886150
  • python3使用concurrent执行多进程任务

    随着计算机技术的发展,诸如GPU和超算平台等越来越发达,这些技术的本质其实并没有带来算法上的革新,之所以能够提升计算的速度和规模,很大程度上是因为分布式和并行计...

    DechinPhy
  • 百度目标检测7日打卡常见问题汇总:PaddleDetection安装,yml配置文件,评估结果等

    环境版本说明: 推荐版本 Paddle:1.8.4 Python: 3.7 CUDA:10.0 pip: 9.0.1+ PaddleDetectio...

    烤粽子
  • 一个简单的更改让PyTorch读取表格数据的速度提高20倍:可大大加快深度学习训练的速度

    在训练深度学习模型时,性能至关重要。数据集可能非常庞大,而低效的训练方法意味着迭代速度变慢,超参数优化的时间更少,部署周期更长以及计算成本更高。

    deephub

扫码关注云+社区

领取腾讯云代金券