请先安装好Anaconda,Python版本为3.6.1。
接着安装好numba
包。
conda install numba
然后,我们需要下载CUDA Toolkit,下载地址:https://developer.nvidia.com/cuda-75-downloads-archive
我用的Mac,所以选择的Mac版本。下载完后,根据提示进行安装。(安装路径,Mac的话,应该在/Developer/NVIDIA/CUDA-7.5/
路径下)
Test
新建文件 test.py
,写入代码:
"""
Python CUDA编程
"""
from __future__ import print_function, division, absolute_import
from timeit import default_timer as timer
from matplotlib.pylab import imshow, jet, show, ion
import numpy as np
from numba import jit
@jit
def mandel(x, y, max_iters):
"""
Given the real and imaginary parts of a complex number,
determine if it is a candidate for membership in the Mandelbrot
set given a fixed number of iterations.
"""
i = 0
c = complex(x,y)
z = 0.0j
for i in range(max_iters):
z = z*z + c
if (z.real*z.real + z.imag*z.imag) >= 4:
return i
return 255
@jit
def create_fractal(min_x, max_x, min_y, max_y, image, iters):
height = image.shape[0]
width = image.shape[1]
pixel_size_x = (max_x - min_x) / width
pixel_size_y = (max_y - min_y) / height
for x in range(width):
real = min_x + x * pixel_size_x
for y in range(height):
imag = min_y + y * pixel_size_y
color = mandel(real, imag, iters)
image[y, x] = color
return image
image = np.zeros((500 * 2, 750 * 2), dtype=np.uint8)
s = timer()
create_fractal(-2.0, 1.0, -1.0, 1.0, image, 20)
e = timer()
print(e - s)
imshow(image)
#jet()
#ion()
show()
将输出图片,并将输出图片创建所需时间:0.36721072000364074
。可以尝试注释掉@jit
,纯粹使用Python生成图片的时间,大概有9-10倍的差距。