使用 Numba 限制活动 CPU 核心的数量

Question

@njit(parallel=True)当使用 Numba 库中的装饰器运行函数时，我试图限制活动 CPU 核心的数量。

到目前为止，我已经做了这样的事情（示例）：

from numba import njit, prange, set_num_threads

import numpy as np


@njit(parallel=True)
def func():

    for i in prange(int(1e10)):

        print(np.cos(i))


if __name__ == '__main__':
    
    num_threads = input(f'Enter max number of threads to use: ')

    set_num_threads(int(num_threads))

    func()

但即使将线程数设置为 1，使用任务管理器进行快速检查显示工作负载仍然由所有核心共享（在我的例子中为 4）。

在下面的屏幕截图中可以看到这一点。红色圆圈表示程序启动和中断的时间。

任务管理器截图

有没有更好的方法来限制活动核心的数量？

Answer 1

您通过使用打印功能并将其显示在控制台中来给您的CPU核心带来压力，而不是因为计算。在此示例中，我在打印之前仅累积一次值。这样，你可以更容易地看到代码的真实cpu使用情况。

from numba import njit, prange, set_num_threads

import numpy as np


@njit(parallel=True)
def func():

    k = 0
    for i in prange(int(1e10)):
        k = k + np.cos(i)

    print(k)



if __name__ == '__main__':

    num_threads = input(f'Enter max number of threads to use: ')

    set_num_threads(int(num_threads))