用点积广播

Question

我正在尝试求解一组点积。目前，我发现最好的方法是使用这个循环。

for i in range(n):
    u[:,:,i] = np.dot(K[i,:,:], v[:,:,i])

有人可以通过广播或其他更快的方法帮助我解决这个问题吗？

非常感谢

Answer 1

一个简单的解决方法是：

u = np.matmul(K, v.transpose(2, 0, 1))
u = u.transpose(1, 2, 0)

我假设 K 的形状为 n x b x c，v 的形状为 c x d x n。像这样转置 v 使其成为 n x c x d 形状，这使得可以对输入形状 b x c 和 c x d 执行 n 次矩阵乘法，结果为 n x b x d。

最后，您还需要转置 u 以获得所需的 b x d x n 形状。

Answer 2

比手动绕轴移动更直观可读的解决方案是使用 np.einsum 来计算结果：

import numpy as np

n = 7
b = 8
c = 11
d = 4

K = np.random.rand(n, b, c)
v = np.random.rand(c, d, n)

u = np.matmul(K, v.transpose(2, 0, 1)).transpose(1, 2, 0)
w = np.einsum('nbc,cdn->bdn', K, v)

print(np.allclose(u, w)) # True

用点积广播

Broadcasting with dot product

numpy

broadcast