python - 使用 parallel=True 时如何防止 Numba 并行化特定循环？

Question

我有一个像这样的测试代码：

import numpy as np
from numba import njit, prange
import random


def func():
    a = np.empty(0)
    for i in range(10):
        b = np.arange(np.random.randint(10))
        a = np.concatenate((a, b))
        print(len(b), ' / ', len(a))


func()

它按预期工作：

3  /  3
6  /  9
5  /  14
8  /  22
9  /  31
6  /  37
9  /  46
5  /  51
3  /  54
9  /  63

如果我通过添加装饰器来使用 Numba 进行编译@njit(parallel=True)，它会停止工作，因为 Numba 会并行化这个导致问题的循环：

2  /  0
1  /  0
2  /  0
7  /  0
0  /  0
3  /  0
2  /  0
5  /  0
7  /  0
2  /  0

在我的真实代码中，此函数中有另一个使用 prange 的循环，因此我希望在整个函数上使用 parallel=True。但我不希望它用于这个循环。

如何防止这种情况？

标签： pythonpython-3.xparallel-processingnumba

Answer 1

我最终将这部分代码从具有的@njit'ed 函数中拆分出来parallel=True，并且在这个被分解的函数上我使用@njit但离开了parallel=False

这为我做到了。