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()
Image size: (512L, 512L)
CPU: 0.66s
kernel config: (16, 64) x (32, 8)
GPU: 0.09s
[Finished in 61.4s]
Numpy提供常见的数学函数,包含许多有用的数学库
Scipy是python下的数值计算库,和Numpy一样是科学计算不可或缺的库
Matplotlib是用以绘制二维图形的Python模块
Numbapro是CUDA提供的专用库
timeit是计时工具