如何加速numpy张量*张量运算

Question

我的代码有一个瓶颈，它是 numpy 3d 数组乘以 * operator 与 numpy 3d 数组。
我想用 numba @njit 或 @jit 装饰器来加速这部分程序，但它降低了 2 倍的性能。
慢的部分代码：

@numba.jit
def mat_mul_and_sum(img1, img2, alpha):
    return img1*(1-alpha) + img2*alpha

img1、img2 和 alpha 是具有相同形状的 3d np.array。
是否可以加快这行代码的速度？

Answer 1

如果按如下方式展开循环，对于大小为 (100, 100, 100) 的数组，numba 的速度是纯 numpy 版本的两倍，这可能是因为不需要分配中间数组:

import numpy as np
import numba as nb

def mat_mul_and_sum(img1, img2, alpha):
    return img1*(1-alpha) + img2*alpha


@nb.jit
def mat_mul_and_sum2(img1, img2, alpha):
    NI, NJ, NK = img1.shape
    out = np.empty((NI, NJ, NK))

    for i in range(NI):
        for j in range(NJ):
            for k in range(NK):
                out[i,j,k] = img1[i,j,k] * (1.0 - alpha[i,j,k]) + img2[i,j,k] * alpha[i,j,k]

    return out

然后测试：

N = 100
img1 = np.random.normal(size=(N, N, N))
img2 = np.random.normal(size=(N, N, N))
alpha = np.random.normal(size=(N, N, N))

A = mat_mul_and_sum(img1, img2, alpha)
B = mat_mul_and_sum2(img1, img2, alpha)

np.allclose(A,B) #True

%timeit mat_mul_and_sum(img1, img2, alpha)
# 4.6 ms ± 44.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit mat_mul_and_sum2(img1, img2, alpha)
# 2.4 ms ± 10.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

更新： 您也可以尝试将装饰器更改为 nb.jit(parallel=True)，然后将外部循环替换为 for i in nb.prange(NI):，这在我的机器上将结果从 timeit 降低到 1.37 毫秒。这个时间和其他时间肯定会因机器而异，也会因输入的大小而异。

Answer 2

一个选项实际上是 numba 按照应有的方式使用（不仅仅是应用装饰器）。但是，对于您的特定功能，您可以使用 numexpr 包来使用多核渲染。

import numexpr as ne

def mat_mul_and_sum_numexpr(a, b, alpha):
    return ne.evaluate('a*(1-alpha) + b*alpha')

使用其他答案的时间：

In [11]: %timeit mat_mul_and_sum(img1, img2, alpha)
21.6 ms ± 955 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [12]: %timeit mat_mul_and_sum2(img1, img2, alpha)
6.35 ms ± 126 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [13]: %timeit mat_mul_and_sum_numexpr(img1, img2, alpha)
4.22 ms ± 54.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [14]: np.allclose(mat_mul_and_sum(img1, img2, alpha), mat_mul_and_sum_numexpr(img1, img2, alpha))
Out[14]: True

您可能能够通过 numba 的并行化获得一些额外的性能，但通常使用 numexpr 可以提供不错的性能提升，而无需重写任何代码。

如何加速numpy张量*张量运算

How to speed up numpy tensor*tensor operation

python

optimization

numpy

numba