Python：3d 与 3d 数组的元素乘法答案

【问题标题】：Python: Element-wise multiplication of 3d with 3d arraysPython：3d 与 3d 数组的元素乘法
【发布时间】：2022-01-01 10:35:37
【问题描述】：

我在使用 Python 以高性能方式实现以下等式时遇到了一些问题：

beta 和 gamma 是笛卡尔坐标 {x,y} 和 b,m 是一些索引值，可以很大 n=10000。对于 l=2 和 m,b = 4 的简单情况（l 和 m 始终具有相同的长度），我有一个工作版本的代码如下所示。我使用 timeit 检查了代码，瓶颈是元素乘法与大小为 (3,3) 的数组以及将结果数组重新整形为形状 (3m,3m)。有人知道如何提高性能吗？（我还注意到我当前的版本对于较大的 l.... 值有相当大的开销）

import numpy as np 

g_l3 = np.array([[1, 4, 5],[2, 6, 7]])
g_l33  = g_l3.reshape(-1, 3, 1) * g_l3.reshape(-1, 1, 3)

A_lm = np.arange(1, 9, 1).reshape(2, 4)
B_lb = np.arange(7, 15, 1).reshape(2, 4)

AB_lmb = A_lm.reshape(-1, 4, 1) * B_lb.reshape(-1, 1, 4)

D_lmb33 = np.sum(g_l33.reshape(-1, 1, 1, 3, 3) * AB_lmb.reshape(-1, 4, 4, 1, 1), axis=0)
D = np.concatenate(np.concatenate(D_lmb33, axis=2), axis=0)

【问题讨论】：

numpy.reshape 几乎不需要时间。它不会对数据进行任何真正的更改，它只是重新解释已经存在的缓冲区。而你的g_l33 不是逐元素乘法——那是在做矩阵乘法。结果是 2x3x3。这是你所期望的吗？
当前代码很快：在我的机器上大约需要 0.025 毫秒。我不理解可能会很慢的情况（我应该帮助提供一个选择的参数导致执行缓慢的情况）。我也不明白(3m,3m)的来源（代码中没有出现）。
所有乘法都是矩阵乘法。你最终会得到一个 (3m)**2 数组。如果m 是任何大小，那将会非常大。
我不确定，但这个公式在我看来就像一个 einsum。 np.einsum('lg,lb,lmb,lxb->xmbg',g_l_gamma,g_l_beta,A_l_m_beta,B_l_b_beta,optimize='optimal') 之类的东西（b 被 x 替换）。但这不是您在代码中所做的...

标签： python numpy performance coding-efficiency

【解决方案1】：

In [387]: %%timeit
     ...: g_l3 = np.array([[1, 4, 5],[2, 6, 7]])
     ...
     ...: D_lmb33 = np.sum(g_l33.reshape(-1, 1, 1, 3, 3) * AB_lmb.reshape(-1, 4,
     ...:  4, 1, 1), axis=0)
     ...: D = np.concatenate(np.concatenate(D_lmb33, axis=2), axis=0)
     ...: 
     ...: 
70.7 µs ± 226 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

检查碎片，并用newaxis 重写reshape，这对我来说视觉上更清晰——尽管速度基本相同：

In [388]: g_l3.shape
Out[388]: (2, 3)
In [389]: g_l33.shape
Out[389]: (2, 3, 3)
In [390]: np.allclose(g_l33, g_l3[:,:,None]*g_l3[:,None,:])
Out[390]: True
In [391]: AB_lmb.shape
Out[391]: (2, 4, 4)
In [392]: np.allclose(AB_lmb, A_lm[:,:,None]*B_lb[:,None,:])
Out[392]: True

所以这些常见的outer 产品在二维数组的最后一维上。

还有一个outer，

In [393]: temp=g_l33.reshape(-1, 1, 1, 3, 3) * AB_lmb.reshape(-1, 4, 4, 1, 1)
In [394]: temp.shape
Out[394]: (2, 4, 4, 3, 3)
In [396]: np.allclose(temp, g_l33[:,None,None,:,:] * AB_lmb[:, :,:, None,None])
Out[396]: True

这些可能可以组合成一个表达式，但这不是必需的。

D_lmb33 对前导维度求和：

In [405]: D_lmb33.shape
Out[405]: (4, 4, 3, 3)

双重连接也可以通过转置和重塑来完成：

In [406]: np.allclose(D_lmb33.transpose(1,2,0,3).reshape(12,12),D)
Out[406]: True

总体而言，您的代码似乎有效地利用了numpy。对于较大的前导维度，(N,4,4,3,3) 中间数组可能很大，并且需要时间。但在numpy 本身内没有其他选择。我认为代数不允许使用更早地进行求和。使用numba 或numexpr 另一个问题。

【讨论】：