概述

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

但是,即使我传入@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,dtype = np.int32)                     

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

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

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

print(end - start)                                        

但是,如果您确实知道可以并行化代码,则可以始终手动使用线程或进程.只是改编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,b)
timefunc(correct,"numba (1 thread)",func_nb,"numba (%d threads)" % nthreads,func_nb_mt,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[:],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,"numba (2 threads)",func_nb_mt2,"numba (3 threads)",func_nb_mt3,"numba (4 threads)",func_nb_mt4,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

总结

以上是编程之家为你收集整理的如何使numba @jit使用所有cpu核心(parallelize numba @jit)全部内容，希望文章能够帮你解决如何使numba @jit使用所有cpu核心(parallelize numba @jit)所遇到的程序开发问题。

如果您也喜欢它,动动您的小指点个赞吧