Python GPU编程之加速篇

Python 2015-12-25

在上一篇中我们对Python GPU编程进行了入门讲解,这里我们将继续探究它的应用。 我们以FFT (离散傅氏变换的快速算法)为例,看看究竟如何利用GPU进行加速。 先看示例代码,然后进行讲解。

import sys
import numpy as np
from scipy.signal import fftconvolve
from scipy import misc, ndimage
from matplotlib import pyplot as plt
from numbapro.cudalib import cufft
from numbapro import cuda, vectorize
from timeit import default_timer as timer

@vectorize(['complex64(complex64, complex64)'], target='gpu')
#目标平台是64位机器且拥有GPU
def vmult(a, b):
    return a * b

def best_grid_size(size, tpb):
    bpg = np.ceil(np.array(size, dtype=np.float) / tpb).astype(np.int).tolist()
    return tuple(bpg)

def main():
    # 构建过滤器
    laplacian_pts = '''
    -4 -1 0 -1 -4
    -1 2 3 2 -1
    0 3 4 3 0
    -1 2 3 2 -1
    -4 -1 0 -1 -4
    '''.split()

    laplacian = np.array(laplacian_pts, dtype=np.float32).reshape(5, 5)

    # 构建图像
    try:
        filename = sys.argv[1]
        image = ndimage.imread(filename, flatten=True).astype(np.float32)
    except IndexError:
        image = misc.lena().astype(np.float32)

    print("Image size: %s" % (image.shape,))

    response = np.zeros_like(image)
    response[:5, :5] = laplacian

    # CPU
    ts = timer()
    cvimage_cpu = fftconvolve(image, laplacian, mode='same')
    te = timer()
    print('CPU: %.2fs' % (te - ts))

    # GPU
    threadperblock = 32, 8
    blockpergrid = best_grid_size(tuple(reversed(image.shape)), threadperblock)
    print('kernel config: %s x %s' % (blockpergrid, threadperblock))

    # cuFFT系统,触发器初始化.
    # 对于小数据集来说,效果更明显.
    # 不应该计算这里浪费的时间
    cufft.FFTPlan(shape=image.shape, itype=np.complex64, otype=np.complex64)

    # 开始GPU运行计时
    ts = timer()
    image_complex = image.astype(np.complex64)
    response_complex = response.astype(np.complex64)

    d_image_complex = cuda.to_device(image_complex)
    d_response_complex = cuda.to_device(response_complex)

    cufft.fft_inplace(d_image_complex)
    cufft.fft_inplace(d_response_complex)

    vmult(d_image_complex, d_response_complex, out=d_image_complex)

    cufft.ifft_inplace(d_image_complex)

    cvimage_gpu = d_image_complex.copy_to_host().real / np.prod(image.shape)

    te = timer()
    print('GPU: %.2fs' % (te - ts))

    # 绘制结果
    plt.subplot(1, 2, 1)
    plt.title('CPU')
    plt.imshow(cvimage_cpu, cmap=plt.cm.gray)
    plt.axis('off')

    plt.subplot(1, 2, 2)
    plt.title('GPU')
    plt.imshow(cvimage_gpu, cmap=plt.cm.gray)
    plt.axis('off')

    plt.show()

if __name__ == '__main__':
    main()

看看运行结果:

1.png

时间对比:

Image size: (512L, 512L)
CPU: 0.66s
kernel config: (16, 64) x (32, 8)
GPU: 0.09s
[Finished in 61.4s]

1.导入的包

  • Numpy提供常见的数学函数,包含许多有用的数学库
  • Scipy是python下的数值计算库,和Numpy一样是科学计算不可或缺的库
  • Matplotlib是用以绘制二维图形的Python模块
  • Numbapro是CUDA提供的专用库
  • timeit是计时工具

2.首先程序提供了一个测试数据集并进行转化,对于CPU部分,直接调用Scipy中的fftconvolve函数计算出结果,而GPU则主要调用了numbapro中cufft库,具体使用参考官方文档。

3.计算出结果后绘图展示,主要使用matplotlib的方法,只需设置plt对象并传入参数即可。

以上在代码中有详细注释,不一样展开。总的来说,使用GPU还是得参考官方的文档,并结合Python编程,才能解决复杂的问题并达到良好的效果。


本文由 Tony 创作,采用 知识共享署名 3.0,可自由转载、引用,但需署名作者且注明文章出处。

如果对您有用,您的支持将鼓励我继续创作!