python - Numba 无法并行化循环
问题描述
以下代码运行,但 prange 循环无法并行化,尽管它显然是可并行化的:
import numpy as np
from numba import njit, prange
@njit(parallel=True)
def f1(money, u, v, cost_u, cost_v):
# task: find index of u iu_opt and index of v iv_opt
# such that u[iu_opt] + v[iv_opt] is maximal subject
# to cost_u[iu_opt] + cost_v[iv_opt] < money
na = money.size
ncu = cost_u.size
ncv = cost_v.size
iu_opt = np.empty((na,),dtype=np.int16)
iv_opt = np.empty((na,),dtype=np.int16)
for ia in prange(na):
money_i = money[ia]
ivbest = 0 # initially pick iv = 0
# find max iu corresponding to iv = 0
for iu in range(ncu-1,-1,-1):
if cost_u[iu] + cost_v[0] < money_i: break
iubest = iu
# compute initial score
score_best = u[iu] + v[0]
# then try to increase iv
for iv in range(1,ncv):
# it not enough money for u_0
if cost_v[iv] + cost_u[0] > money_i: break
while cost_v[iv] + cost_u[iu] > money_i:
iu -= 1
assert iu >= 0
score_now = u[iu] + v[iv]
if score_now > score_best:
ivbest = iv
iubest = iu
score_best = score_now
iu_opt[ia] = iubest
iv_opt[ia] = ivbest
return iu_opt, iv_opt
na = 50
ncu = 204
ncv = 205
money = np.arange(na)/(na)
cost_u = np.arange(ncu)/ncu
u = np.cumsum(np.random.random_sample(ncu))
cost_v = np.arange(ncv)/ncv
v = np.cumsum(np.random.random_sample(ncv))
iu, iv = f1(money, u, v, cost_u, cost_v)
f1.parallel_diagnostics(level=4)
如果这有帮助,那么问题的设置如下:对于每个值 money[ ia] 我们想要找到 i 和 j 最大化 u[i] + v[j],因为 cost_u[i] + cost_v[j] 不能超过 money[ia]。
解决方案
如果有人遇到过类似的问题,我最终通过将大循环的内部拆分为另一个 njit 函数来使其工作。这是代码:
import numpy as np
from numba import njit, prange
@njit(parallel=True)
def f1(money, u, v, cost_u, cost_v):
# task: find index of u iu_opt and index of v iv_opt
# such that u[iu_opt] + v[iv_opt] is maximal subject
# to cost_u[iu_opt] + cost_v[iv_opt] < money
na = money.size
iu_opt = np.empty((na,),dtype=np.int16)
iv_opt = np.empty((na,),dtype=np.int16)
for ia in prange(na):
money_i = money[ia]
iubest, ivbest = f1_int(money_i,u,v,cost_u,cost_v)
iu_opt[ia] = iubest
iv_opt[ia] = ivbest
return iu_opt, iv_opt
@njit
def f1_int(money_i,u,v,cost_u,cost_v):
ivbest = 0 # initially pick iv = 0
ncu = cost_u.size
ncv = cost_v.size
# find max iu corresponding to iv = 0
for iu in range(ncu-1,-1,-1):
if cost_u[iu] + cost_v[0] < money_i: break
iubest = iu
# compute initial score
score_best = u[iu] + v[0]
# then try to increase iv
for iv in range(1,ncv):
# it not enough money for u_0
if cost_v[iv] + cost_u[0] > money_i: break
while cost_v[iv] + cost_u[iu] > money_i:
iu -= 1
assert iu >= 0
score_now = u[iu] + v[iv]
if score_now > score_best:
ivbest = iv
iubest = iu
score_best = score_now
return iubest, ivbest
na = 50
ncu = 204
ncv = 205
money = np.arange(na)/(na)
cost_u = np.arange(ncu)/ncu
u = np.cumsum(np.random.random_sample(ncu))
cost_v = np.arange(ncv)/ncv
v = np.cumsum(np.random.random_sample(ncv))
iu, iv = f1(money, u, v, cost_u, cost_v)
f1.parallel_diagnostics(level=4)
这并不能真正回答为什么会发生问题的问题,但会以某种方式起作用。
推荐阅读
- html - Css网格,如何在不弄乱整个页面的情况下调整元素的大小
- firebase - Flutter Web/Firebase - 在浏览器中按回绕过验证过程
- c# - Windows 窗体设计器未加载
- javascript - 使用可选的 SFC 自定义 VueJS 应用程序
- python - 如何避免重叠的误差线并替换 X 和 Y 轴?
- python - 如何在python中逐字阅读?
- python - 如何解决 optimizer.step() 不起作用的问题
- spring-boot - 如何将@PathVariable 值绑定到Spring Boot Rest API 中的@Where 子句
- swift - 属性类型“MyViewModel”与其包装器类型的“wrappedValue”属性不匹配
- redis - 阻塞直到完成计数