簡單的numba + CUDA 實測
起因
numba + CUDA
numba天生支持NumPy，但是CUDA部分僅提供非常有限的支持
CUDA部分代碼
簡單的numba + CUDA 實測
起因
一時興起，是我太閑了吧。
最近需要對一個4k圖像進行單個像素級別的處理，由於用python用得人有點懶，直接上python在所有像素上循環一遍。每個像素做的工作其實很簡單，就是判斷一下這個像素是否符合某一准則，如果符合就將這個像素mask上。噼里啪啦寫好一個腳本，一運行居然等了很久，都沒有結果。一開還以為是不是哪里寫錯了，進入了無限循環什么的，但是最后發現其實執行效率就是那么低。我做了一個實例，在一個3008x4112像素的圖像上進行簡單的分類。（實例中的兩個分類功能都可以用cv2直接實現，這里僅作實例進行測試使用。）
from __future__ import print_function
import cv2
import math
import numpy as np
from numpy.linalg import norm
import time
H = 3008
W = 4112
class Validator(object):
def __init__(self):
pass
def is_valid(self, x, y):
return False
class RadiusValidator(Validator):
def __init__(self, center, R, width):
super(RadiusValidator, self).__init__()
self.R = R
self.center = center # A two element NumPy array. [x, y].
self.width = width
if ( self.width <= 0 ):
raise Exception("self.width wrong. self.width = {}".format(self.width))
# Overide parent's function.
def is_valid(self, x, y):
x = x - self.center[0]
y = y - self.center[1]
r = math.sqrt( x * x + y * y )
if ( r >= self.R - self.width and r <= self.R + self.width ):
return True
else:
return False
class PolarLineSegmentValidator(Validator):
def __init__(self, center, theta, length, width):
super(PolarLineSegmentValidator, self).__init__()
self.center = center # A two element NumPy array. [x, y].
self.theta = theta
self.length = length
self.width = width
self.endP = np.zeros((2,), dtype=np.float32)
self.endP[0] = self.length * math.cos( self.theta )
self.endP[1] = self.length * math.sin( self.theta )
self.dir = self.endP / self.length
def is_valid(self, x, y):
# Transform the coordinate to the local frame.
v = np.zeros((2, ), dtype=np.float32)
v[0] = x - self.center[0]
v[1] = y - self.center[1]
# Projection.
proj = self.dir.dot( v )
if ( proj < 0 ):
return False
elif ( proj > self.length ):
return False
# Othogonal vector.
oth = v - proj*self.dir
# Distance.
d = norm( oth )
if ( d <= self.width ):
return True
else:
return False
if __name__ == "__main__":
print("List the indices for AEAG brightness evaluation.")
# Create a blank image with all zeros.
img = np.zeros((H, W), dtype=np.uint8)
imgCenter = np.array([ W/2, H/2 ], dtype=np.float32)
# Create the validators.
vList = []
vList.append( RadiusValidator( imgCenter, min( W/2, H/2 ) * 0.75, 20 ) )
vList.append( PolarLineSegmentValidator( imgCenter, math.pi/2, min( W/2, H/2 ) * 0.75, 10 ) )
vList.append( PolarLineSegmentValidator( imgCenter, math.pi/4, min( W/2, H/2 ) * 0.75, 10 ) )
vList.append( PolarLineSegmentValidator( imgCenter, 0, min( W/2, H/2 ) * 0.75, 10 ) )
vList.append( PolarLineSegmentValidator( imgCenter, -math.pi/4, min( W/2, H/2 ) * 0.75, 10 ) )
vList.append( PolarLineSegmentValidator( imgCenter, -math.pi/2, min( W/2, H/2 ) * 0.75, 10 ) )
# Record the starting time.
start = time.time()
for i in range(H):
for j in range(W):
flag = False
for v in vList:
flag = flag or v.is_valid( j, i )
if ( True == flag ):
img[i, j] = 255
if ( 0 == i % 100 ):
print("%d," % (i), end=None)
# Record the ending time.
end = time.time()
print(end - start)
# Save the image.
cv2.imwrite("ValidPixels.png", img)
上面的腳本在我的機器上執行時間約520s，圖像像素約為1200萬。之前早有耳聞說是像我這樣簡單的進行for循環，python的效率會非常低，今天算是見識到了。於是想着，得，還是改成C++的吧。但是轉念一想，貌似有一個叫什么“牛B”的庫可以用，可以使用CUDA進行加速。簡單查了一下果然是有這個庫，numba。
可能是我太閑，我決定試一下。
numba + CUDA
numba可以在沒有CUDA支持時使用CPU進行加速，而這里我只感興趣CUDA的部分。
numba要用conda配置，也是醉了。還好用了conda environment。
我想說numba的文檔風格我有點不適應，也許是我看的太粗略，一時間沒有參透其中的道理。官方的reference manual之外，還有一個github上的項目對numba的CUDA部分進行了解釋。這個項目是gtc2018-numba。
於是我趟了幾個坑，這里總結一下。
numba天生支持NumPy，但是CUDA部分僅提供非常有限的支持
在使用numba時，主要數據類型與NumPy非常接近，大部分NumPy的屬性和接口都有對應的實現，有些接口的參數數量和種類可能受到限制，這需要仔細看numba的官方文檔關於NumPy的一部分。但是中的一個坑是，numba的CUDA部分對python的支持更受限制，這需要參考官方文檔的CUDA部分。很容易進的坑包括
kernel函數內部對大部分NumPy的函數不支持。
CUDA kernel函數內不允許任何可能產生memory allocation的NumPy函數，這使得絕大部分NumPy函數都不能使用了，尤其是創建array的函數，例如zeros(), array()，和linspace()等。對array進行運算的函數也不支持。但是我們仍然可以使用NumPy的標量函數，例如sin()。kernel函數內部可以使用NumPy對象的shape屬性。NumPy的indexing和slicing操作似乎仍然可以使用。
kernel函數內部對list貌似不支持，對tuple支持。貌似可以創建list，但是索引和更新時會失敗。一個在stackoverflow上的討論。dictionary和generator就不要想了。
kernel函數內部貌似不支持對自定義類的使用。
device函數需要使用@cuda.jit(device=True)進行修飾，可以有返回值。我不知道為什么官方和上述github教程對device函數沒有專門的章節，也許是我沒有找到。
於是過了一些遠大於520s的時間以后，如下的測試代碼終於可以運行了。
CUDA部分代碼
from __future__ import print_function
import cv2
import math
import matplotlib.pyplot as plt
import numpy as np
from numba import cuda
import time
@cuda.jit(device=True)
def d_radius_validate(cx, cy, R, width, x, y):
x = x - cx
y = y - cy
r = math.sqrt( x * x + y * y )
if ( r >= R - width and r <= R + width ):
return 255
else:
return 0
@cuda.jit(device=True)
def d_polar_line_segment_validate(cx, cy, theta, length, width, x, y):
n0 = length * math.cos(theta) / length
n1 = length * math.sin(theta) / length
v0 = x - cx
v1 = y - cy
proj = n0 * v0 + n1 * v1
if ( proj < 0 ):
return 0
elif ( proj > length ):
return 0
oth0 = v0 - proj*n0
oth1 = v1 - proj*n1
d = math.sqrt( oth0 * oth0 + oth1 * oth1 )
if ( d <= width ):
return 255
else:
return 0
@cuda.jit
def k_validate(imgOut):
tx = cuda.blockIdx.x*cuda.blockDim.x + cuda.threadIdx.x
ty = cuda.blockIdx.y*cuda.blockDim.y + cuda.threadIdx.y
xStride = cuda.blockDim.x * cuda.gridDim.x
yStride = cuda.blockDim.y * cuda.gridDim.y
cx, cy = 2056, 1504
R = 1504*0.75
halfWidth = 10.0
for y in range( ty, imgOut.shape[0], yStride ):
for x in range( tx, imgOut.shape[1], xStride ):
flag = d_radius_validate( cx, cy, R, halfWidth, 1.0*x, 1.0*y )
if ( 0 != flag ):
imgOut[y, x] = flag
flag = d_polar_line_segment_validate( cx, cy, -math.pi/2, R, halfWidth, 1.0*x, 1.0*y )
if ( 0 != flag ):
imgOut[y, x] = flag
flag = d_polar_line_segment_validate( cx, cy, -math.pi/4, R, halfWidth, 1.0*x, 1.0*y )
if ( 0 != flag ):
imgOut[y, x] = flag
flag = d_polar_line_segment_validate( cx-halfWidth, cy, 0.0, R + halfWidth, halfWidth, 1.0*x, 1.0*y )
if ( 0 != flag ):
imgOut[y, x] = flag
flag = d_polar_line_segment_validate( cx, cy, math.pi/4, R, halfWidth, 1.0*x, 1.0*y )
if ( 0 != flag ):
imgOut[y, x] = flag
flag = d_polar_line_segment_validate( cx, cy, math.pi/2, R, halfWidth, 1.0*x, 1.0*y )
if ( 0 != flag ):
imgOut[y, x] = flag
if __name__ == "__main__":
# Create an image.
img = np.zeros((3008, 4112), dtype=np.int32)
# Record the starting time.
start = time.time()
dImg = cuda.to_device(img)
cuda.synchronize()
k_validate[[100, 100, 1], [16, 16, 1]](dImg)
cuda.synchronize()
img = dImg.copy_to_host()
# Record the ending time.
end = time.time()
print(end - start)
# Save the image.
cv2.imwrite("ValidPixels_numba.png", img)
print("Done.")
如上述代碼所示，采用100x100的grid，和16x16的block thread配置，在當前的GPU（GeForce GTX 980Ti，拿不上台面啊）上的執行時間約為0.4s。
嗯，快了。
也許我真的是閑的。
————————————————
版權聲明：本文為CSDN博主「風海流」的原創文章，遵循 CC 4.0 BY-SA 版權協議，轉載請附上原文出處鏈接及本聲明。
原文鏈接：https://blog.csdn.net/huyaoyu/article/details/89742577