专栏首页Dechin的专栏Python的GPU编程实例——近邻表计算

Python的GPU编程实例——近邻表计算

技术背景

GPU加速是现代工业各种场景中非常常用的一种技术,这得益于GPU计算的高度并行化。在Python中存在有多种GPU并行优化的解决方案,包括之前的博客中提到的cupy、pycuda和numba.cuda,都是GPU加速的标志性Python库。这里我们重点推numba.cuda这一解决方案,因为cupy的优势在于实现好了的众多的函数,在算法实现的灵活性上还比较欠缺;而pycuda虽然提供了很好的灵活性和相当高的性能,但是这要求我们必须在Python的代码中插入C代码,这显然是非常不Pythonic的解决方案。因此我们可以选择numba.cuda这一解决方案,只要在Python函数前方加一个numba.cuda.jit的修饰器,就可以在Python中用最Python的编程语法,实现GPU的加速效果。

加速场景

我们需要先了解的是,GPU在什么样的计算场景下能够实现加速的效果,很显然的是,并不是所有的计算过程都能在GPU上表现出加速的效果。前面说道,GPU的加速作用,是源自于高度的并行化,所谓的并行,就要求进程之前互不干扰或者依赖。如果说一个进程的计算过程或者结果,依赖于另一个进程中的计算结果,那么就无法实现完全的并行,只能使用串行的技术。这里为了展示GPU加速的效果,我们就引入一个在分子动力学模拟领域中常见的问题:近邻表的计算。

近邻表计算的问题是这样描述的:给定一堆数量为n的原子系统,每一个原子的三维坐标都是已知的,给定一个截断常数

d_0

,当两个原子之间的距离

d_{i,j}<=d_0

时,则认为这两个原子是相邻近的原子。那么最终我们需要给出一个0-1矩阵

A_{i,j}

,当

A_{i,j}=0

时,表示

i,j

两个原子互不相邻,反之则相邻。那么对于这个问题场景,我们就可以并行化的遍历

n\times n

的空间,直接输出

A_{n\times n}

大小的近邻表。这个计算场景是一个非常适合用GPU来加速的计算,以下我们先看一下不用GPU加速时的常规实现方案:

# cuda_neighbor_list.py

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

@jit
def neighbor_list(crd, neighbors, data_length, cutoff):
    """CPU based neighbor list calculation.
    """
    for i in range(data_length):
        for j in range(i+1, data_length):
            if np.linalg.norm(crd[i]-crd[j]) <= cutoff:
                neighbors[i][j] = 1
                neighbors[j][i] = 1
    return neighbors

if __name__ == '__main__':
    np.random.seed(1)
    atoms = 2**2
    cutoff = 0.5
    crd = np.random.random((atoms,3))
    adjacent = np.zeros((atoms, atoms))
    adjacent = neighbor_list(crd, adjacent, atoms, cutoff)
    print (adjacent)

这是最常规的一种CPU上的实现方案,遍历所有的原子,计算原子间距,然后填充近邻表。这里我们还使用到了numba.jit即时编译的功能,这个功能是在执行到相关函数时再对其进行编译的方法,在矢量化的计算中有可能使用到芯片厂商所提供的SIMD的一些优化。当然,这里都是CPU层面的执行和优化,执行结果如下:

$ python3 cuda_neighbor_list.py 
[[0. 0. 0. 0.]
 [0. 0. 1. 0.]
 [0. 1. 0. 1.]
 [0. 0. 1. 0.]]

这个输出的结果就是一个0-1近邻表。

基于Numba的GPU加速

对于上述的近邻表计算的场景,我们很容易的想到这个neighbor_list函数可以用GPU的函数来进行改造。对于每一个

d_{i,j}

我们都可以启动一个线程去执行计算,类似于CPU上的SIMD技术,GPU中的这项优化称为SIMT。而在Python中改造成GPU函数的方法也非常简单,只需要把函数前的修饰器改一下,去掉函数内部的for循环,就基本完成了,比如下面这个改造的近邻表计算的案例:

# cuda_neighbor_list.py

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

@jit
def neighbor_list(crd, neighbors, data_length, cutoff):
    """CPU based neighbor list calculation.
    """
    for i in range(data_length):
        for j in range(i+1, data_length):
            if np.linalg.norm(crd[i]-crd[j]) <= cutoff:
                neighbors[i][j] = 1
                neighbors[j][i] = 1
    return neighbors

@cuda.jit
def cuda_neighbor_list(crd, neighbors, cutoff):
    """GPU based neighbor list calculation.
    """
    i, j = cuda.grid(2)
    dis = ((crd[i][0]-crd[j][0])**2+\
           (crd[i][1]-crd[j][1])**2+\
           (crd[i][2]-crd[j][2])**2)**0.5
    neighbors[i][j] = dis <= cutoff[0] and dis > 0

if __name__ == '__main__':
    import time
    np.random.seed(1)

    atoms = 2**5
    cutoff = 0.5
    cutoff_cuda = cuda.to_device(np.array([cutoff]).astype(np.float32))
    crd = np.random.random((atoms,3)).astype(np.float32)
    crd_cuda = cuda.to_device(crd)
    adjacent = np.zeros((atoms, atoms)).astype(np.float32)
    adjacent_cuda = cuda.to_device(adjacent)

    time0 = time.time()
    adjacent_c = neighbor_list(crd, adjacent, atoms, cutoff)
    time1 = time.time()
    cuda_neighbor_list[(atoms, atoms), (1, 1)](crd_cuda,
                                               adjacent_cuda,
                                               cutoff_cuda)
    time2 = time.time()
    adjacent_g = adjacent_cuda.copy_to_host()
    print ('The time cost of CPU with numba.jit is: {}s'.format(\
                                            time1-time0))
    print ('The time cost of GPU with cuda.jit is: {}s'.format(\
                                            time2-time1))
    print ('The result error is: {}'.format(np.sum(adjacent_c-\
                                            adjacent_g)))

需要说明的是,当前Numba并未支持所有的numpy的函数,因此有一些计算的功能需要我们自己去手动实现一下,比如计算一个Norm的值。这里我们在输出结果中不仅统计了结果的正确性,也给出了运行的时间:

$ python3 cuda_neighbor_list.py 
The time cost of CPU with numba.jit is: 0.6401469707489014s
The time cost of GPU with cuda.jit is: 0.19208502769470215s
The result error is: 0.0

需要说明的是,这里仅仅运行了一次的程序,而jit即时编译的加速效果在第一次的运行中其实并不明显,甚至还有一些速度偏慢,但是在后续过程的函数调用中,就能够起到比较大的加速效果。所以这里的运行时间并没有太大的代表性,比较有代表性的时间对比可以看如下的案例:

# cuda_neighbor_list.py

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

@jit
def neighbor_list(crd, neighbors, data_length, cutoff):
    """CPU based neighbor list calculation.
    """
    for i in range(data_length):
        for j in range(i+1, data_length):
            if np.linalg.norm(crd[i]-crd[j]) <= cutoff:
                neighbors[i][j] = 1
                neighbors[j][i] = 1
    return neighbors

@cuda.jit
def cuda_neighbor_list(crd, neighbors, cutoff):
    """GPU based neighbor list calculation.
    """
    i, j = cuda.grid(2)
    dis = ((crd[i][0]-crd[j][0])**2+\
           (crd[i][1]-crd[j][1])**2+\
           (crd[i][2]-crd[j][2])**2)**0.5
    neighbors[i][j] = dis <= cutoff[0] and dis > 0

if __name__ == '__main__':
    import time
    np.random.seed(1)

    atoms = 2**10
    cutoff = 0.5
    cutoff_cuda = cuda.to_device(np.array([cutoff]).astype(np.float32))
    crd = np.random.random((atoms,3)).astype(np.float32)
    crd_cuda = cuda.to_device(crd)
    adjacent = np.zeros((atoms, atoms)).astype(np.float32)
    adjacent_cuda = cuda.to_device(adjacent)
    time_c = 0.0
    time_g = 0.0

    for _ in range(100):
        time0 = time.time()
        adjacent_c = neighbor_list(crd, adjacent, atoms, cutoff)
        time1 = time.time()
        cuda_neighbor_list[(atoms, atoms), (1, 1)](crd_cuda,
                                                adjacent_cuda,
                                                cutoff_cuda)
        time2 = time.time()
        if _ != 0:
            time_c += time1 - time0
            time_g += time2 - time1
    
    print ('The total time cost of CPU with numba.jit is: {}s'.format(\
                                            time_c))
    print ('The total time cost of GPU with cuda.jit is: {}s'.format(\
                                            time_g))

这个案例中也没有修改较多的地方,只是把一次计算的时间调整为多次计算的时间,并且忽略第一次计算过程中的即时编译,最终输出结果如下:

$ python3 cuda_neighbor_list.py 
The total time cost of CPU with numba.jit is: 14.955506563186646s
The total time cost of GPU with cuda.jit is: 0.018685102462768555s

可以看到,在GPU加速后,相比于CPU的高性能运算,能够提速达将近1000倍!

总结概要

对于Pythoner而言,苦其性能已久。如果能够用一种非常Pythonic的方法来实现GPU的加速效果,对于Pythoner而言无疑是巨大的好消息,Numba就为我们提供了这样的一个基础功能。本文通过一个近邻表计算的案例,给出了适用于GPU加速的计算场景。这种计算场景可并行化的程度较高,而且函数会被多次用到(在分子动力学模拟的过程中,每一个step都会调用到这个函数),因此这是一种最典型的、最适用于GPU加速场景的案例。

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

我来说两句

0 条评论
登录 后参与评论

相关文章

  • 学界 | Facebook AI实验室开源相似性搜索库Faiss:性能高于理论峰值55%,提速8.5倍

    在用户日常搜索过程中,一个经常出现的问题即大多数返回的网站结果拥有完全相同或者几乎一样的信息。而应用了相似性搜索的相似引擎即可为用户返回最恰当、最合适的结果,同...

    AI科技评论
  • LeCun推荐:最新PyTorch图神经网络库,速度快15倍(GitHub+论文)

    过去十年来,深度学习方法(例如卷积神经网络和递归神经网络)在许多领域取得了前所未有的成就,例如计算机视觉和语音识别。

    磐创AI
  • GitHub 上 57 款最流行的开源深度学习项目

    本文整理了 GitHub 上最流行的 57 款深度学习项目(按 stars 排名)。最后更新:2016.08.09 1.TensorFlow 使用数据流图计算可...

    陆勤_数据人网
  • GitHub 上 57 款最流行的开源深度学习项目

    1.TensorFlow 使用数据流图计算可扩展机器学习问题 TensorFlow 是谷歌的第二代机器学习系统,按照谷歌所说,在某些基准测试中,TensorF...

    顶级程序员
  • Pytorch 分布式模式介绍

    数据较多或者模型较大时,为提高机器学习模型训练效率,一般采用多GPU的分布式训练。

    狼啸风云
  • python machine learning package

    如果您目前正在使用Python进行机器学习项目,那么您可能已经听说过这个流行的开源库,称为Tensorflow。该库是由谷歌与Brain Team合作开发的。T...

    乐说科技
  • QA派|GNN工业应用-PinSAGE

    Pinterest是一个图片素材网站,pins是指图片,而boards则是图片收藏夹的意思。

    Houye
  • 用 GPU 加速 TSNE:从几小时到几秒

    原标题 | Accelerating TSNE with GPUs: From hours to seconds

    AI研习社
  • Python手写机器学习最简单的KNN算法

    今天开始,我打算写写机器学习教程。说实话,相比爬虫,掌握机器学习更实用竞争力也更强些。

    AI科技大本营
  • 【GNN】PinSAGE:GCN 在工业级推荐系统中的应用

    今天学习的是 Pinterest 和斯坦福大学 2018 年合作的论文《Graph Convolutional Neural Networks for Web-...

    阿泽 Crz
  • 如何玩转谷歌TensorFlow? | 牛人讲堂

    AI并不是一门简单的学科,AI算法的开发和调试并没有一个统一的、集成了大量API方便调用的平台和语言,目前的人工智能开发平台仍然处于一种半蛮荒的状态。许多功能需...

    AI科技评论
  • 机器学习各语言领域工具库中文版汇总

    我在鹅厂做安全
  • PyTorch入门笔记-基本数据类型

    本小节主要介绍 PyTorch 中的基本数据类型,先来看看 Python 和 PyTorch 中基本数据类型的对比。

    触摸壹缕阳光
  • 机器之心论文解读:可用于十亿级实时检索的循环二分嵌入模型(RBE)

    论文链接:https://arxiv.org/pdf/1802.06466.pdf

    机器之心
  • Python3实现打格点算法的GPU加速

    在数学和物理学领域,总是充满了各种连续的函数模型。而当我们用现代计算机的技术去处理这些问题的时候,事实上是无法直接处理连续模型的,绝大多数的情况下都要转化成一个...

    DechinPhy
  • 上手必备!不可错过的TensorFlow、PyTorch和Keras样例资源

    TensorFlow、Keras和PyTorch是目前深度学习的主要框架,也是入门深度学习必须掌握的三大框架,但是官方文档相对内容较多,初学者往往无从下手。本人...

    AI科技大本营
  • 比DGL快14倍:PyTorch图神经网络库PyG上线了

    项目链接:https://github.com/rusty1s/pytorch_geometric

    机器之心
  • 比DGL快14倍:PyTorch图神经网络库PyG上线了

    项目链接:https://github.com/rusty1s/pytorch_geometric

    刀刀老高
  • 比DGL快14倍:PyTorch图神经网络库PyG上线了

    项目链接:https://github.com/rusty1s/pytorch_geometric

    昱良

扫码关注云+社区

领取腾讯云代金券