python - Numba 没有提高性能

Question

我正在测试一些采用 numpy 数组的函数的 numba 性能，并比较:

import numpy as np
from numba import jit, vectorize, float64
import time
from numba.core.errors import NumbaWarning
import warnings

warnings.simplefilter('ignore', category=NumbaWarning)

@jit(nopython=True, boundscheck=False) # Set "nopython" mode for best performance, equivalent to @njit
def go_fast(a):     # Function is compiled to machine code when called the first time
    trace = 0.0
    for i in range(a.shape[0]):   # Numba likes loops
        trace += np.tanh(a[i, i]) # Numba likes NumPy functions
    return a + trace              # Numba likes NumPy broadcasting
   
class Main(object):
    def __init__(self) -> None:
        super().__init__()
        self.mat     = np.arange(100000000, dtype=np.float64).reshape(10000, 10000)

    def my_run(self):
        st = time.time()
        trace = 0.0
        for i in range(self.mat.shape[0]):   
            trace += np.tanh(self.mat[i, i]) 
        res = self.mat + trace
        print('Python Diration: ', time.time() - st)
        return res                           
    
    def jit_run(self):
        st = time.time()
        res = go_fast(self.mat)
        print('Jit Diration: ', time.time() - st)
        return res
        
obj = Main()
x1 = obj.my_run()
x2 = obj.jit_run()

输出是:

Python Diration:  0.2164750099182129
Jit Diration:  0.5367801189422607

如何获得此示例的增强版本？

Answer 1

Numba 实现的执行时间较慢是由于编译时间，因为 Numba 在使用时编译函数(只有第一次，除非参数类型发生变化)。它这样做是因为它无法在调用函数之前知道参数的类型。希望您可以为 Numba 指定参数类型，以便它可以直接编译函数(当执行装饰器函数时)。这是结果代码:

@njit('float64[:,:](float64[:,:])')
def go_fast(a):
    trace = 0.0
    for i in range(a.shape[0]):
        trace += np.tanh(a[i, i])
    return a + trace

请注意，njit 是 jit+nopython=True 的快捷方式，并且 boundscheck 已设置为False 默认情况下(参见 doc )。

在我的机器上，这导致 Numpy 和 Numba 的执行时间相同。实际上，执行时间不受 tanh 函数计算的限制。它受表达式 a + trace 的限制(对于 Numba 和 Numpy)。预计执行时间相同，因为两者的实现方式相同:它们创建一个临时的新数组来执行加法。由于 page faults，创建一个新的临时数组很昂贵以及 RAM 的使用(a 完全从 RAM 中读取，临时数组完全存储在 RAM 中)。如果你想要更快的计算，那么你需要执行操作就地(这可以防止页面错误和 x86 平台上昂贵的 cache-line write allocations)。