在 3D 矩阵的情况下用 numpy 操作替换 for 循环答案

【问题标题】：Replacing for loops with numpy operations in case of 3D matrices在 3D 矩阵的情况下用 numpy 操作替换 for 循环
【发布时间】：2021-12-31 14:22:35
【问题描述】：

我有这个数学公式要实现 ![https://i.stack.imgur.com/S28BA.png] 其中例如 w_fk 表示形状矩阵 (F, K)。我已将其实现为

gamma_dashed_lft = np.zeros((L, F, T))
for l in range(L):
    for f in range(F):
        for t in range(T):
            temp = 0
            for k in range(K):
                temp = temp + (q_lk[l, k] * w_fk[f, k] * h_kt[k, t])
            gamma_dashed_lft[l, f, t] = temp
return gamma_dashed_lft

在给定公式的情况下，用矩阵乘法替换 for 循环的方法是什么？

【问题讨论】：

标签： python numpy matrix matrix-multiplication

【解决方案1】：

你应该提供一个具体的例子。幸运的是，读取尺寸并创建并不难：

In [302]: L,F,T,K=2,3,4,5
In [303]: q_lk=np.arange(L*K).reshape(L,K)
In [304]: w_fk=np.arange(F*K).reshape(F,K)
In [305]: h_kt=np.arange(K*T).reshape(K,T)

应用于您的代码时会产生：

In [306]: gamma_dashed_lft = np.zeros((L, F, T))
     ...: for l in range(L):
     ...:     for f in range(F):
     ...:         for t in range(T):
     ...:             temp = 0
     ...:             for k in range(K):
     ...:                 temp = temp + (q_lk[l, k] * w_fk[f, k] * h_kt[k, t])
     ...:             gamma_dashed_lft[l, f, t] = temp
     ...: 

In [308]: gamma_dashed_lft
Out[308]: 
array([[[ 400.,  430.,  460.,  490.],
        [1000., 1080., 1160., 1240.],
        [1600., 1730., 1860., 1990.]],

       [[1000., 1080., 1160., 1240.],
        [2600., 2855., 3110., 3365.],
        [4200., 4630., 5060., 5490.]]])

充分利用broadcasting的等价表达式是：

In [309]: arr =(q_lk[:,None,:,None]*w_fk[None,:,:,None]*h_kt[None,None,:,:]).sum(axis=2)
In [310]: arr.shape
Out[310]: (2, 3, 4)
In [311]: np.allclose(arr,gamma_dashed_lft)
Out[311]: True

在设置广播时，我的目标是一个形状为 (L,F,K,T) 且在 K 上求和的数组。

既然你让我创建测试用例，我让你制定广播细节。这对你来说将是一个很好的锻炼。

einsum

In [446]: D=np.einsum('lk,fk,kt->lft', q_lk, w_fk, h_kt)
In [447]: D.shape
Out[447]: (2, 3, 4)
In [448]: arr =(q_lk[:,None,:,None]*w_fk[None,:,:,None]*h_kt[None,None,:,:]).sum
     ...: (axis=2)
In [449]: np.allclose(arr,D)
Out[449]: True
In [450]: timeit arr =(q_lk[:,None,:,None]*w_fk[None,:,:,None]*h_kt[None,None,:,
     ...: :]).sum(axis=2)
22.4 µs ± 2.02 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [451]: timeit D=np.einsum('lk,fk,kt->lft', q_lk, w_fk, h_kt)
12.2 µs ± 40.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

matmul

In [458]: timeit E=((q_lk[:,None,None,:]*w_fk[None,:,None,:])@h_kt[None,None,:,:,: ]).squeeze()
10.6 µs ± 44.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

为此，我使用 (1,1,K,T) 将 (L,F,1,K) 数组添加到 @，从而得到 (L,F,1,T)。 LF 是 matmul 'batch' 维度，而 K 是总和维度。

【讨论】：

这是如何工作的？ q_lk 等是二维的，但第 309 行中的表达式似乎对它们进行了 4 维索引。我也想知道这个的效率。与dot（从 BLAS 调用 DGEMM）的矩阵-矩阵乘法比基于一维数组的运算（例如内积或外积）快 20 倍。
所有那些 None 在索引中添加维度。如果您对此不熟悉，请查找 np.newaxis。就像q_lk.reshape(L,1,K,1)。在求和之前，乘积是 (2,3,5,4)。 numpy 中的中间数组可能很大。这是使用已编译构建块的部分成本。
如果问题可以转换为矩阵乘积，matmul/dot 相当快，因为它可以传递给 BLAS 函数。您的示例确实有一个 temp 总和，但也有 3 个产品，所以我没有尝试制定任何类似 dot 的格式。我刚刚测试了一个更快的einsum 替代方案。

【解决方案2】：

产品q[l,k]*w[f,k] 会针对t 的每个值进行评估，所以我会这样处理：

for l,k:
  qw[k] = q[l,k] * w[f,k]

然后在内部循环中使用该临时数组qw。现在您有一堆内积，实际上是矩阵向量积，并且您在运算中节省了 T 的因子。

但是，这并没有像矩阵-矩阵乘法那样为您提供缓存优化。为此，为每个l 创建一个矩阵qw_l：

for l:
  qw_l[f,k] = q[l,k]*w[f,k]
  matrix-matrix-product: qw_l times h

（请注意，我没有把所有的循环都拼出来！）

这会花费您一个大小为F x K 的数组，但这可能不是问题。如果您想知道，创建它的成本是 F x K，而矩阵矩阵乘积是 F x T x K，所以您不必担心。

【讨论】：