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

Question

我有这个数学公式要实现 ![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 循环的方法是什么？

Answer 1

乘积 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，所以您不必担心。

Answer 2

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

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)

为此，我正在制作一个 (L,F,1,K) 数组到 @ 和 (1,1,K,T)，结果是 (L,F,1,吨）。 LF 是 matmul 'batch' 维度，而 K 是乘积维度。

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

Replacing for loops with numpy operations in case of 3D matrices

python

numpy

matrix

matrix-multiplication

einsum

matmul