使用numba cuda 加速Python运算
习惯了cuda c,可能会认为cuda和c才是黄金档搭。
Python作为一种开发效率比较高的脚本语言,有助于我们快速实现某种功能。
但是Python执行效率极其之慢。
这种情况下,用cuda的高并发特性,来提升Python执行速度,是一种很好的选择。
1.随机数生成
随机数生成是一项很重要的功能。
当Python自带的random,np.random在cuda函数中无法直接使用时,这是一个非常头疼的事。
有一个方法是将随机数/序列提前在cuda函数外实现好,再传递给cuda核函数使用,但是这就要占用cuda的显存,同时还要考虑加载数据的时间。
幸好的事numba提供了numba.cuda.random
,可以便于我们生成随机数。
numba random官方网页中提供了一个示例,通过均匀分布来实现pi的计算。
由于numba.cuda.random.xoroshiro128p_normal_float64
默认生成
N
(
0
,
1
)
N(0,1)
N(0,1)分布序列。
这里提供一个使用numba.cuda.random
来生成复合高斯分布(如均值为100,方差为30的)的随机数:
有 N ( μ , s i g m a ) N(\mu,sigma) N(μ,sigma)分布的序列转成 N ( 0 , 1 ) N(0,1) N(0,1),标准化公式为:
y = x − μ δ \qquad\qquad y=\cfrac{x-\mu}{\sqrt{\delta}} y=δx−μ
故从有 N ( 0 , 1 ) N(0,1) N(0,1)分布的序列转成 N ( μ , s i g m a ) N(\mu,sigma) N(μ,sigma)分布,为:
y = δ ⋅ x + μ \qquad\qquad y=\sqrt{\delta} \cdot x+\mu y=δ⋅x+μ
代码如下:
from numba import cuda
from numba.cuda.random import create_xoroshiro128p_states, xoroshiro128p_normal_float64
import numpy as np
import math
@cuda.jit
def random_gen(rng_states, out):
"""Find the maximum value in values and store in result[0]"""
thread_id = cuda.grid(1)
print("thread_id",thread_id)
out[thread_id]=xoroshiro128p_normal_float64(rng_states, thread_id)
out[thread_id]=int(out[thread_id]*math.sqrt(30)+100)
threads_per_block = 16
blocks = 16
rng_states = create_xoroshiro128p_states(threads_per_block * blocks, seed=1)
out = np.zeros((threads_per_block * blocks), dtype=np.float32)
out_d = cuda.to_device(out)
random_gen[blocks, threads_per_block](rng_states, out_d)
out = out_d.copy_to_host()
print('\n', out)
产生如下序列:
[ 92. 100. 97. 101. 95. 103. 101. 105. 92. 101. 100. 97. 91. 90.
97. 104. 100. 98. 97. 102. ...]
用numpy可求得均值和方差分别为:
99.609375 30.902099609375
生成整数随机序列,可以通过均匀分布,再经过适当放缩、平移实现,如采用(0,1)均匀分布实现[0,100]整数的均匀采样:
int(100*xoroshiro128p_uniform_float64(rng_states, col))
注意:rng_states 中的个数不宜设置过大,否则也可能会大大影响gpu运行速度。
2.numba njit与numba cuda性能的比较
以二维卷积运算为例。关于卷积运算,可参考离散信号的卷积与相关。
对大小为1111*1111
的方阵进行卷积运算(卷积核为11*11
)。
import numpy as np
import time
k=1111
spot=np.random.randint(0,10,(k,k)).astype(np.float64)
kernel=np.random.randint(0,10,(11,11)).astype(np.float64)
2.1 scipy实现
from scipy import signal
t1=time.time()
spot_conv = signal.convolve2d(spot, kernel, mode='same')
t2=time.time()
print(t2-t1)
耗时约为0.31s。
2.2 个人编程实现
所用代码同离散信号的卷积与相关。
def Convolve2D(img, kernel,pad=None):
kernel=np.flipud(np.fliplr(kernel)) #卷积核反转
h_k,w_k = kernel.shape
h_img,w_img = img.shape
if pad is None:
pad=(h_k-1)//2
if pad != 0:
img_pad = np.zeros((img.shape[0] + pad*2, img.shape[1] + pad*2),dtype=np.float64)
img_pad[int(pad):int(-1 * pad), int(pad):int(-1 * pad)] = img
else:
img_pad = img
out = np.zeros((h_img, w_img),dtype=np.float64)
for y in prange(img.shape[1]):
for x in prange(img.shape[0]):
out[x, y] = np.sum(kernel * img_pad[x: x + h_k, y: y + w_k])
return out
t3=time.time()
spot_conv2=Convolve2D(spot, kernel)
t4=time.time()
print(t4-t3)
耗时约为2s。
2.3 numba njit实现
import numba
from numba import njit,prange
@numba.njit(parallel=True)
def Convolve2D(img, kernel,pad=None):
kernel=np.flipud(np.fliplr(kernel)) #卷积核反转
h_k,w_k = kernel.shape
h_img,w_img = img.shape
if pad is None:
pad=(h_k-1)//2
if pad != 0:
img_pad = np.zeros((img.shape[0] + pad*2, img.shape[1] + pad*2),dtype=np.float64)
img_pad[int(pad):int(-1 * pad), int(pad):int(-1 * pad)] = img
else:
img_pad = img
out = np.zeros((h_img, w_img),dtype=np.float64)
for y in prange(img.shape[1]):
for x in prange(img.shape[0]):
out[x, y] = np.sum(kernel * img_pad[x: x + h_k, y: y + w_k])
return out
t3=time.time()
for i in range(10):
spot_conv2=Convolve2D(spot, kernel)
t4=time.time()
print((t4-t3)/10)
耗时约为0.07s。已经相当不错了。
这里不要把prange
写成range
,否则会报错:Segmentation fault。
2.4 numba cuda实现
import numba
from numba import cuda
@cuda.jit
def conv_cuda(img,img_conv,kernel,h,w, ks):
threadId = cuda.grid(1)
row, col = threadId//w, threadId % w
if (row < h and col < w):
res = 0.0
for i in range(ks):
for j in range(ks):
curRow = row - ks // 2 + i
curCol = col - ks // 2 + j
if (0 <= curRow < h and 0 <= curCol < w):
value = img[curRow*w+curCol]
res += value * kernel[i*ks+j]
cuda.syncthreads()
img_conv[ row*w+col] = res
#img_conv变量是必须的,
#如果直接将值赋给img[row*w+col],由于同步问题其他线程就有可能使用的是卷积后的数,
#而不是原始数据
h,w=img.shape
img_d = cuda.to_device(np.ascontiguousarray(img.reshape(-1)))
img_conv_d = cuda.to_device(np.ascontiguousarray(np.zeros(h*w)))
kernel=np.flipud(np.fliplr(kernel))
kernel_d = cuda.to_device(np.ascontiguousarray(kernel.reshape(-1)))
ks=kernel.shape[0]
t5=time.time()
nthreads = 1024
nblocks = (h*w // nthreads) + 1
for i in range(10):
conv_cuda[nblocks, nthreads](img_d,img_conv_d,kernel_d,h,w, ks)
img_conv3=img_conv_d.copy_to_host()
img_conv3=img_conv3.reshape(h,w)
t6=time.time()
print((t6-t5)/10)
最终结果为0.03s。
变量在从host复制到device前,最好使用np.ascontiguousarray
转为连续存储。
如果是多张图片进行卷积运算,设置设置二维的线程数。
@cuda.jit
def imgs_conv_cuda(imgs,imgs_conv, imgs_num, psf, h,w,ks):
imgId= cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
imgPixel_Id =cuda.threadIdx.y + cuda.blockIdx.y * cuda.blockDim.y
#cuda.grid(2)
row, col = imgPixel_Id//w, imgPixel_Id % w
if imgId < imgs_num:
if (row < h and col < w):
intens = 0.0
for i in range(ks):
for j in range(ks):
curRow = row - ks // 2 + i
curCol = col - ks // 2 + j
if (0 <= curRow < h and 0 <= curCol < w):
value = imgs[imgId, curRow*w+curCol]
intens += value * psf[i*ks+j]
imgs_conv[imgId, row*w+col] = intens
获取线程两个维度id号的两种方法:
imgId= cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
imgPixel_Id =cuda.threadIdx.y + cuda.blockIdx.y * cuda.blockDim.y
imgId,imgPixel_Id =cuda.grid(2)
参考文献
[1] https://numba.readthedocs.io/en/stable/
[2] 基于 Numba 的 CUDA Python 编程简介
[3] https://numba.pydata.org/numba-doc/latest/cuda/random.html
[4] An introduction to CUDA in Python (Part 3)