时间戳

首页 > 解决方案 > 如何加快 numpy tensor*tensor 运算

python - 如何加快 numpy tensor*tensor 运算

问题描述

我的代码中有一个瓶颈，它是 numpy 3d 数组乘以* operatornumpy 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。
是否可以加快这行代码的速度？

标签： pythonnumpyoptimizationnumba

解决方案

一种选择实际上是使用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您可能能够通过s 并行化挤出一些额外的性能，但通常使用numexpr可以提供很好的性能提升，而无需重新编写任何代码。

推荐阅读