如何在属于 jitclass 成员的函数中使用 prange?
例如,我想使用 prange 在函数 do_smth
中并行执行。
@jitclass([('x', float64[:])])
class Test:
def __init__(self):
self.x = np.arange(10)
def do_smth(self):
for i in range(len(self.x)):
pass
最佳答案
我还没有找到在 jitclass
成员中指定 parallel=True
的方法。
但是您可以改为调用外部函数。如果确实需要并行化,这个小开销应该是最小的。
@nb.experimental.jitclass([('x', nb.float64[:, :])])
class Test:
def __init__(self):
self.x = np.arange(10000.).reshape((10,1000))
def do_smth(self):
rows, cols = self.x.shape
sums = np.zeros(cols)
for r in nb.prange(rows):
print(r)
sums[r] = np.sum(self.x[r])
return np.sum(sums)
def do_other(self):
return do_other(self.x)
@nb.njit([nb.float64(nb.float64[:,:])], parallel=True)
def do_other(x):
rows, cols = x.shape
sums = np.zeros(cols)
for r in nb.prange(rows):
print(r)
sums[r] = np.sum(x[r])
return np.sum(sums)
t = Test()
print("Class member")
print(t.do_smth())
print("External function")
print(t.do_other())
即使 prange
在 Test.do_smth()
中显式使用,循环也不会并行化。处理行的顺序很好地指示了循环何时并行化:
Class member
0
1
2
3
4
5
6
7
8
9
49995000.0
External function
0
2
6
3
5
9
8
7
4
1
49995000.0
关于python - 如何在jitclass中使用prange,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/67577756/