如何使 numba @jit 使用所有 cpu 内核(并行化 numba @jit)

     2023-02-16     125

关键词:

【中文标题】如何使 numba @jit 使用所有 cpu 内核(并行化 numba @jit)【英文标题】:How to make numba @jit use all cpu cores (parallelize numba @jit) 【发布时间】:2018-01-18 12:01:37 【问题描述】:

我正在使用 numbas @jit 装饰器在 python 中添加两个 numpy 数组。如果我使用@jitpython 相比,性能是如此之高。

但是,即使我传入@numba.jit(nopython = True, parallel = True, nogil = True),它也没有利用所有 CPU 内核

有什么方法可以使用 numba @jit 的所有 CPU 内核。

这是我的代码:

import time                                                
import numpy as np                                         
import numba                                               

SIZE = 2147483648 * 6                                      

a = np.full(SIZE, 1, dtype = np.int32)                     

b = np.full(SIZE, 1, dtype = np.int32)                     

c = np.ndarray(SIZE, dtype = np.int32)                     

@numba.jit(nopython = True, parallel = True, nogil = True) 
def add(a, b, c):                                          
    for i in range(SIZE):                                  
        c[i] = a[i] + b[i]                                 

start = time.time()                                        
add(a, b, c)                                               
end = time.time()                                          

print(end - start)                                        

【问题讨论】:

您的示例不使用prange - 我找到了您的问题,因为我想知道为什么prange 不能并行运行,但我在@jit 中缺少parallel = True.. . 【参考方案1】:

您可以将parallel=True 传递给任何 numba jitted 函数,但这并不意味着它总是使用所有内核。您必须了解 numba 使用一些启发式方法来使代码并行执行,有时这些启发式方法根本无法在代码中找到要并行化的任何内容。当前有一个pull request,因此如果无法使其“并行”,它会发出警告。所以它更像是一个“请尽可能使其并行执行”参数而不是“强制并行执行”。

但是,如果您确实知道可以并行化代码,则始终可以手动使用线程或进程。只是适应example of using multi-threading from the numba docs:

#!/usr/bin/env python
from __future__ import print_function, division, absolute_import

import math
import threading
from timeit import repeat

import numpy as np
from numba import jit

nthreads = 4
size = 10**7  # CHANGED

# CHANGED
def func_np(a, b):
    """
    Control function using Numpy.
    """
    return a + b

# CHANGED
@jit('void(double[:], double[:], double[:])', nopython=True, nogil=True)
def inner_func_nb(result, a, b):
    """
    Function under test.
    """
    for i in range(len(result)):
        result[i] = a[i] + b[i]

def timefunc(correct, s, func, *args, **kwargs):
    """
    Benchmark *func* and print out its runtime.
    """
    print(s.ljust(20), end=" ")
    # Make sure the function is compiled before we start the benchmark
    res = func(*args, **kwargs)
    if correct is not None:
        assert np.allclose(res, correct), (res, correct)
    # time it
    print(':>5.0f ms'.format(min(repeat(lambda: func(*args, **kwargs),
                                          number=5, repeat=2)) * 1000))
    return res

def make_singlethread(inner_func):
    """
    Run the given function inside a single thread.
    """
    def func(*args):
        length = len(args[0])
        result = np.empty(length, dtype=np.float64)
        inner_func(result, *args)
        return result
    return func

def make_multithread(inner_func, numthreads):
    """
    Run the given function inside *numthreads* threads, splitting its
    arguments into equal-sized chunks.
    """
    def func_mt(*args):
        length = len(args[0])
        result = np.empty(length, dtype=np.float64)
        args = (result,) + args
        chunklen = (length + numthreads - 1) // numthreads
        # Create argument tuples for each input chunk
        chunks = [[arg[i * chunklen:(i + 1) * chunklen] for arg in args]
                  for i in range(numthreads)]
        # Spawn one thread per chunk
        threads = [threading.Thread(target=inner_func, args=chunk)
                   for chunk in chunks]
        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join()
        return result
    return func_mt


func_nb = make_singlethread(inner_func_nb)
func_nb_mt = make_multithread(inner_func_nb, nthreads)

a = np.random.rand(size)
b = np.random.rand(size)

correct = timefunc(None, "numpy (1 thread)", func_np, a, b)
timefunc(correct, "numba (1 thread)", func_nb, a, b)
timefunc(correct, "numba (%d threads)" % nthreads, func_nb_mt, a, b)

我突出显示了我更改的部分,其他所有内容都是从示例中逐字复制的。这利用了我机器上的所有内核(4 核机器因此 4 个线程),但没有显示出显着的加速:

numpy (1 thread)       539 ms
numba (1 thread)       536 ms
numba (4 threads)      442 ms

在这种情况下,多线程缺乏(很多)加速是因为加法是一种带宽受限的操作。这意味着从数组中加载元素并将结果放入结果数组中要比实际加法花费更多的时间。

在这些情况下,您甚至会因为并行执行而看到速度变慢!

只有当函数更复杂并且与加载和存储数组元素相比实际操作需要大量时间时,您才会看到并行执行的巨大改进。 numba 文档中的示例是这样的:

def func_np(a, b):
    """
    Control function using Numpy.
    """
    return np.exp(2.1 * a + 3.2 * b)

@jit('void(double[:], double[:], double[:])', nopython=True, nogil=True)
def inner_func_nb(result, a, b):
    """
    Function under test.
    """
    for i in range(len(result)):
        result[i] = math.exp(2.1 * a[i] + 3.2 * b[i])

这实际上(几乎)与线程数成比例,因为两次乘法、一次加法和一次调用math.exp 比加载和存储结果要慢得多:

func_nb = make_singlethread(inner_func_nb)
func_nb_mt2 = make_multithread(inner_func_nb, 2)
func_nb_mt3 = make_multithread(inner_func_nb, 3)
func_nb_mt4 = make_multithread(inner_func_nb, 4)

a = np.random.rand(size)
b = np.random.rand(size)

correct = timefunc(None, "numpy (1 thread)", func_np, a, b)
timefunc(correct, "numba (1 thread)", func_nb, a, b)
timefunc(correct, "numba (2 threads)", func_nb_mt2, a, b)
timefunc(correct, "numba (3 threads)", func_nb_mt3, a, b)
timefunc(correct, "numba (4 threads)", func_nb_mt4, a, b)

结果:

numpy (1 thread)      3422 ms
numba (1 thread)      2959 ms
numba (2 threads)     1555 ms
numba (3 threads)     1080 ms
numba (4 threads)      797 ms

【讨论】:

【参考方案2】:

为了完整起见,在 2018 年(numba v 0.39)你可以这样做

from numba import prange

并在您的原始函数定义中将range 替换为prange,就是这样。

这立即使 CPU 利用率达到 100%,在我的情况下,运行时间从 2.9 秒加快到 1.7 秒(对于 SIZE = 2147483648 * 1,在 16 核 32 线程的机器上)。

更复杂的内核通常可以通过传入fastmath=True 来加快速度。

【讨论】:

如何在一个简单的功能上 Cythonize / 允许 numba.jit:(在网络中查找三角形)

】如何在一个简单的功能上Cythonize/允许numba.jit:(在网络中查找三角形)【英文标题】:HowtoCythonize/allownumba.jitonasimplefunction:(FindingTrianglesinnetwork)【发布时间】:2021-09-1507:46:59【问题描述】:背景故事:我一直在寻找一种高性能... 查看详情

numba 中的@jit 和@vectorize 有啥区别?

...发布时间】:2018-05-1315:00:06【问题描述】:什么时候应该使用@vectorize?我尝试了@jit并显示了下面的那部分代码,fromnumbaimportjit@jitdefkma(g,tem 查看详情

numba 中的 jit 和 autojit 有啥区别?

】numba中的jit和autojit有啥区别?【英文标题】:Whatisthedifferencebetweenjitandautojitinnumba?numba中的jit和autojit有什么区别?【发布时间】:2015-02-1014:17:18【问题描述】:我对@9​​87654322@和autojit之间的区别感到困惑。我读过这个:http://... 查看详情

如何在类的成员函数上使用 numba?

】如何在类的成员函数上使用numba?【英文标题】:HowdoIusenumbaonamemberfunctionofaclass?【发布时间】:2017-06-0518:20:18【问题描述】:我正在使用Numba0.30.1的稳定版本。我可以这样做:importnumbaasnb@nb.jit("void(f8[:])",nopython=True)defcomplicated(x... 查看详情

numba - guvectorize 几乎比 jit 快

】numba-guvectorize几乎比jit快【英文标题】:numba-guvectorizebarelyfasterthanjit【发布时间】:2017-06-0719:47:03【问题描述】:我试图并行化在许多独立数据集上运行的蒙特卡洛模拟。我发现numba的并行guvectorize实现仅比numbajit实现快30-40%。... 查看详情

调用 numba jit 函数时,cProfile 会增加大量开销

】调用numbajit函数时,cProfile会增加大量开销【英文标题】:cProfileaddssignificantoverheadwhencallingnumbajitfunctions【发布时间】:2018-12-2507:40:09【问题描述】:将纯Python无操作函数与以@numba.jit修饰的无操作函数进行比较,即:importnumba@nu... 查看详情

如何制作一个虚拟的无操作 @jit 装饰器?

】如何制作一个虚拟的无操作@jit装饰器?【英文标题】:HowdoImakeadummydo-nothing@jitdecorator?【发布时间】:2020-01-0612:39:03【问题描述】:我希望numba成为一个可选依赖项,这样如果安装它就快,否则就慢。所以当没有安装numba时,我... 查看详情

使用带有 numba njit 功能的字典

...:2019-07-3109:30:11【问题描述】:当输入和返回是字典时,如何使用numba加速函数?我熟悉将numba用于接受数字并返回数组的函数,如下所示:@numba.jit(\'float64[:](int32,int32)\',nopython=True)deff(a,b):#return 查看详情

多个输出和 numba 签名

...问题描述】:也许这很简单,但我想知道当有多个输出时如何在jit装饰器中编写签名。例如:importnumbaasnb@nb.jit([\'???(int32,int32,float(:,:),float(:,:))\'],nopython=True)deffoo(nx,ny,a,b) 查看详情

numba安装和使用

numba是针对python加速的包,类似cython,pypy,优势是代码改动少首先要安装llvmliteapt-getinstallllvm-3.8LLVM_CONFIG=/usr/local/llvm38/3.8.1/lib/llvm-3.8/bin/llvm-configpipinstallllvmlite#看自己路径在哪程序里importnumba@numba.jit()装饰要加速的函数 查看详情

Numba:@jit 中的并行标志在我的带有 numpy 2D 数组的代码中不起作用

】Numba:@jit中的并行标志在我的带有numpy2D数组的代码中不起作用【英文标题】:Numba:parallelflagin@jitdoesnotworkinmycodewithnumpy2Darrays【发布时间】:2022-01-1703:48:59【问题描述】:首先,我问了这个问题Numba:Whyguvectorizeissoslow?。当我发现... 查看详情

与 Python+Numba LLVM/JIT 编译的代码相比,Julia 的性能

】与Python+NumbaLLVM/JIT编译的代码相比,Julia的性能【英文标题】:JuliaperformancecomparedtoPython+NumbaLLVM/JIT-compiledcode【发布时间】:2015-06-1510:15:03【问题描述】:到目前为止我看到的Julia的性能基准,例如在http://julialang.org/,将Julia与... 查看详情

Numba 中的稀疏矩阵

...nNumba【发布时间】:2013-10-2513:21:20【问题描述】:我希望使用Numba(http://numba.pydata.org/)加速我的机器学习算法(用Python编写)。请注意,该算法将稀疏矩阵作为其输入数据。在我的纯Python实现中,我使用了来自Scipy的csr_matrix和相... 查看详情

强化学习技巧五:numba提速python程序(代码片段)

...运算),其运行速度可以接近C或FORTRAN语言。numba使用情况使用numpy数组做大量科学计算时使用for循环时1.numba使用导入numpy、numba及其编译器importnumpyasnpimportnumbafro 查看详情

对于纯 numpy 代码,使用 numba 的收益在哪里?

】对于纯numpy代码,使用numba的收益在哪里?【英文标题】:Wherearethegainsusingnumbacomingfromforpurenumpycode?【发布时间】:2017-11-2821:01:37【问题描述】:我想了解使用Numba在for循环中加速纯numpy代码时的收益来自哪里。是否有任何分析... 查看详情

使用 numba 无法获得与 numpy 元素矩阵乘法相同的值

】使用numba无法获得与numpy元素矩阵乘法相同的值【英文标题】:Can\'tgetsamevaluesasnumpyelementwisematrixmultiplicationusingnumba【发布时间】:2018-03-2218:01:20【问题描述】:我一直在玩numba并尝试实现一个简单的元素矩阵乘法。使用“vectoriz... 查看详情

即使对于巨型矩阵,NUMBA CUDA 也比并行 CPU 慢

...时间】:2020-11-1319:01:02【问题描述】:网上只有少数几个使用cuda进行numba的示例,我发现它们都比并行CPU方法慢。带有CUDA目标和模板的矢量化甚至更糟,所以我尝试创建一个自定义内核。您随处可见的一篇博文是https://gist.gith 查看详情

如何安装和导入openmp通过numba使用?

】如何安装和导入openmp通过numba使用?【英文标题】:Howtoinstallandimportopenmptouseitthroughnumba?【发布时间】:2022-01-1908:42:07【问题描述】:我正在python中使用numba构建程序,我需要使用openmp作为numba线程层。我正在努力让它启动并运... 查看详情