关于Python的GPU编程实例近邻表计算的讲解

目录
  • 技术背景
  • 加速场景
  • 基于Numba的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的原子系统,每一个原子的三维坐标都是已知的,给定一个截断常数d0,当两个原子之间的距离di,j<=d0时,则认为这两个原子是相邻近的原子。那么最终我们需要给出一个0-1矩阵Ai,j,当Ai,j=0时,表示i,j两个原子互不相邻,反之则相邻。那么对于这个问题场景,我们就可以并行化的遍历n×n的空间,直接输出An×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的函数来进行改造。对于每一个di,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加速场景的案例。

以上就是关于Python的GPU编程实例近邻表计算的讲解的详细内容,更多关于Python GPU编程实例的资料请关注我们其它相关文章!

(0)

相关推荐

  • 详解python中GPU版本的opencv常用方法介绍

    引言 本篇是以python的视角介绍相关的函数还有自我使用中的一些问题,本想在这篇之前总结一下opencv编译的全过程,但遇到了太多坑,暂时不太想回看做过的笔记,所以这里主要总结python下GPU版本的opencv. 主要函数说明 threshold():二值化,但要指定设定阈值 blendLinear():两幅图片的线形混合 calcHist() createBoxFilter ():创建一个规范化的2D框过滤器 canny边缘检测 createGaussianFilter():创建一个Ga

  • 运行tensorflow python程序,限制对GPU和CPU的占用操作

    一般情况下,运行tensorflow时,默认会占用可以看见的所有GPU,那么就会导致其它用户或程序无GPU可用,那么就需要限制程序对GPU的占用.并且,一般我们的程序也用不了所有的GPU资源,只是强行霸占着,大部分资源都不会用到,也不会提升运行速度. 使用nvidia-smi可以查看本机的GPU使用情况,如下图,这里可以看出,本机的GPU型号是K80,共有两个K80,四块可用(一个K80包括两块K40). 1.如果是只需要用某一块或某几块GPU,可以在运行程序时,利用如下命令运行:CUDA_VI

  • 关于Python的GPU编程实例近邻表计算的讲解

    目录 技术背景 加速场景 基于Numba的GPU加速 总结概要 技术背景 GPU加速是现代工业各种场景中非常常用的一种技术,这得益于GPU计算的高度并行化.在Python中存在有多种GPU并行优化的解决方案,包括之前的博客中提到的cupy.pycuda和numba.cuda,都是GPU加速的标志性Python库.这里我们重点推numba.cuda这一解决方案,因为cupy的优势在于实现好了的众多的函数,在算法实现的灵活性上还比较欠缺:而pycuda虽然提供了很好的灵活性和相当高的性能,但是这要求

  • python交互式图形编程实例(三)

    本文实例为大家分享了python交互式图形编程实例的第三部代码,供大家参考,具体内容如下 #!/usr/bin/env python3 # -*- coding: utf-8 -*- #时钟 from turtle import * from datetime import * def Skip(step): penup() forward(step) pendown() def mkHand(name, length): #注册Turtle形状,建立表针Turtle reset() Skip(

  • python交互式图形编程实例(二)

    本文实例为大家分享了python交互式图形编程的第二部分代码,供大家参考,具体内容如下 #!/usr/bin/env python3 # -*- coding: utf-8 -*- #画个笑脸 from graphics import * win = GraphWin() face = Circle(Point(100,95), 50) leftEye = Circle(Point(80,80) , 5) leftEye.setFill("yellow") leftEye.setOut

  • python交互式图形编程实例(一)

    本文实例为大家分享了python交互式图形编程的具体代码,供大家参考,具体内容如下 #!/usr/bin/env python3# -*- coding: utf-8 -*- #温度转换 from graphics import * win = GraphWin("摄氏温度转换器", 400, 300) win.setCoords(0.0, 0.0, 3.0, 4.0) # 绘制接口 Text(Point(1,3), " 摄氏温度:").draw(win) Text

  • Python threading多线程编程实例

    Python 的多线程有两种实现方法: 函数,线程类 1.函数 调用 thread 模块中的 start_new_thread() 函数来创建线程,以线程函数的形式告诉线程该做什么 复制代码 代码如下: # -*- coding: utf-8 -*- import thread def f(name):   #定义线程函数   print "this is " + name   if __name__ == '__main__':   thread.start_new_thread(f

  • 从零学python系列之数据处理编程实例(二)

    在上一节从零学python系列之数据处理编程实例(一)的基础上数据发生了变化,文件中除了学生的成绩外,新增了学生姓名和出生年月的信息,因此将要成变成:分别根据姓名输出每个学生的无重复的前三个最好成绩和出生年月 数据准备:分别建立四个文本文件 james2.txt     James Lee,2002-3-14,2-34,3:21,2.34,2.45,3.01,2:01,2:01,3:10,2-22 julie2.txt        Julie Jones,2002-8-17,2.59,2.11

  • python并发和异步编程实例

    关于并发.并行.同步阻塞.异步非阻塞.线程.进程.协程等这些概念,单纯通过文字恐怕很难有比较深刻的理解,本文就通过代码一步步实现这些并发和异步编程,并进行比较.解释器方面本文选择python3,毕竟python3才是python的未来,并且python3用原生的库实现协程已经非常方便了. 1.准备阶段 下面为所有测试代码所需要的包 #! python3 # coding:utf-8 import socket from concurrent import futures from selecto

  • python socket通信编程实现文件上传代码实例

    这篇文章主要介绍了python socket通信编程实现文件上传代码实例,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下 写一个file_receive.py和一个file_send.py程序,由file_send.py上传一个文件,file_receive.py接收上传的文件,写到指定的包内 #file_receive.py import socket,subprocess,os BASE_DIR = os.path.dirname(os.pa

  • Python函数式编程实例详解

    本文实例讲述了Python函数式编程.分享给大家供大家参考,具体如下: 函数式编程就是一种抽象程度很高的编程范式,从计算机硬件->汇编语言->C语言->Python抽象程度越高.越贴近于计算,但执行效率也越低.纯粹的函数式编程语言编写的函数没有变量,因此,任意一个函数,只要输入是确定的,输出就是确定的,这种纯函数我们称之为没有副作用.而允许使用变量的程序设计语言,由于函数内部的变量状态不确定,同样的输入,可能得到不同的输出,因此,这种函数是有副作用的.函数式编程的一个特点就是,允许把函数

  • Python原始套接字编程实例解析

    这篇文章主要介绍了Python原始套接字编程实例解析,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下 在实验中需要自己构造单独的HTTP数据报文,而使用SOCK_STREAM进行发送数据包,需要进行完整的TCP交互. 因此想使用原始套接字进行编程,直接构造数据包,并在IP层进行发送,即采用SOCK_RAW进行数据发送. 使用SOCK_RAW的优势是,可以对数据包进行完整的修改,可以处理IP层上的所有数据包,对各字段进行修改,而不受UDP和TCP的

随机推荐