numba njit 的混合数据类型输入

Question

我有一个大数组用于操作，例如矩阵转置。 numba 更快：

#test_transpose.py
import numpy as np
import numba as nb
import time


@nb.njit('float64[:,:](float64[:,:])', parallel=True)
def transpose(x):
    r, c = x.shape
    x2 = np.zeros((c, r))
    for i in nb.prange(c):
        for j in range(r):
            x2[i, j] = x[j][i]
   return x2


if __name__ == "__main__":
    x = np.random.randn(int(3e6), 50)
    t = time.time()
    x = x.transpose().copy()
    print(f"numpy transpose: {round(time.time() - t, 4)} secs")

    x = np.random.randn(int(3e6), 50)
    t = time.time()
    x = transpose(x)
    print(f"numba paralleled transpose: {round(time.time() - t, 4)} secs")

在 Windows 命令提示符中运行

D:\data\test>python test_transpose.py
numpy transpose: 2.0961 secs
numba paralleled transpose: 0.8584 secs

但是，我想输入另一个大矩阵，它们是整数，使用 x 因为

x = np.random.randint(int(3e6), size=(int(3e6), 50), dtype=np.int64)

引发异常，因为

Traceback (most recent call last):
  File "test_transpose.py", line 39, in <module>
    x = transpose(x)
  File "C:\Program Files\Python38\lib\site-packages\numba\core\dispatcher.py", line 703, in _explain_matching_error
    raise TypeError(msg)
TypeError: No matching definition for argument type(s) array(int64, 2d, C)

它无法识别输入数据矩阵为整数。如果我释放整数矩阵的数据类型检查，因为

@nb.njit(parallel=True) # 'float64[:,:](float64[:,:])'
def transpose(x):
    r, c = x.shape
    x2 = np.zeros((c, r))
    for i in nb.prange(c):
        for j in range(r):
            x2[i, j] = x[j][i]
    return x2

它速度较慢：

D:\Data\test>python test_transpose.py
numba paralleled transpose: 1.6653 secs

使用 @nb.njit('int64[:,:](int64[:,:])', parallel=True)正如预期的那样，整数数据矩阵更快。

那么，我如何仍然允许混合数据类型输入但保持速度，而不是为不同类型分别创建函数？

Answer 1

那么，我如何仍然允许混合数据类型输入但保持速度，而不是为不同类型分别创建函数？

问题在于 Numba 函数仅针对 float64 类型定义，而不是针对 int64 类型。类型规范是必需的，因为 Numba 将 Python 代码编译为具有明确定义的类型的本机代码。您可以向 Numba 函数添​​加多个签名：

@nb.njit(['float64[:,:](float64[:,:])', 'int64[:,:](int64[:,:])'], parallel=True)
def transpose(x):
    r, c = x.shape
    # Specifying the dtype is very important here.
    # This is a good habit to take to avoid numerical issues and 
    # slower performance in Numpy too.
    x2 = np.zeros((c, r), dtype=x.dtype)
    for i in nb.prange(c):
        for j in range(r):
            x2[i, j] = x[j][i]
   return x2

比较慢

这是因为延迟编译。第一次执行包括编译时间。当指定签名时，情况并非如此，因为使用了即时编译。

numba 更快

考虑到使用了很多核心，所以这里没有太多。事实上，朴素转置在大矩阵上效率非常低（在这种情况下，在大数组上浪费了大约 90% 的内存吞吐量）。有更快的算法。有关更多信息，请阅读这篇文章 （它只考虑就地 2D 平方转置，这要简单得多，但想法是相同的）。另请注意，类型越宽，数组越大。数组越大，转置速度越慢。