用 numpy tensordot 进行张量乘法

Question

我有一个由n个(d,k)维矩阵和一个(k,n)维矩阵V组成的张量U。

我想将它们相乘，以便结果 returns 维度为 (d,n) 的矩阵，其中第 j 列是 U 的矩阵 j 与第 j 列之间的矩阵乘法结果的 V.

获得此信息的一种可能方法是：

for j in range(n):
    res[:,j] = U[:,:,j] * V[:,j]

我想知道是否有使用 numpy 库的更快方法。特别是我在考虑 np.tensordot() 函数。

这个小片段允许我将单个矩阵乘以一个标量，但是对向量的明显概括并没有返回我所希望的结果。

a = np.array(range(1, 17))
a.shape = (4,4)
b = np.array((1,2,3,4,5,6,7))
r1 = np.tensordot(b,a, axes=0)

有什么建议吗？

Answer 1

有几种方法可以做到这一点。首先想到的是np.einsum:

# some fake data
gen = np.random.RandomState(0)
ni, nj, nk = 10, 20, 100
U = gen.randn(ni, nj, nk)
V = gen.randn(nj, nk)

res1 = np.zeros((ni, nk))
for k in range(nk):
    res1[:,k] = U[:,:,k].dot(V[:,k])

res2 = np.einsum('ijk,jk->ik', U, V)

print(np.allclose(res1, res2))
# True

np.einsum用Einstein notation表示张量收缩。在上面的表达式'ijk,jk->ik'中，i、j和k是分别对应U和V不同维度的下标。每个以逗号分隔的分组对应于传递给 np.einsum 的操作数之一（在本例中 U 的维度为 ijk，V 的维度为 jk）。 '->ik' 部分指定输出数组的维度。输出字符串中不存在的任何带有下标的维度都会被求和。

np.einsum 对于执行复杂的张量收缩非常有用，但可能需要一段时间才能完全了解它的工作原理。您应该查看文档中的示例（上面链接）。

一些其他选项：

与broadcasting逐元素相乘，然后求和：
```
res3 = (U * V[None, ...]).sum(1)
```

inner1d 带有大量移调：

from numpy.core.umath_tests import inner1d

res4 = inner1d(U.transpose(0, 2, 1), V.T)

一些基准：

In [1]: ni, nj, nk = 100, 200, 1000

In [2]: %%timeit U = gen.randn(ni, nj, nk); V = gen.randn(nj, nk)
   ....: np.einsum('ijk,jk->ik', U, V)
   ....: 
10 loops, best of 3: 23.4 ms per loop

In [3]: %%timeit U = gen.randn(ni, nj, nk); V = gen.randn(nj, nk)
(U * V[None, ...]).sum(1)
   ....: 
10 loops, best of 3: 59.7 ms per loop

In [4]: %%timeit U = gen.randn(ni, nj, nk); V = gen.randn(nj, nk)
inner1d(U.transpose(0, 2, 1), V.T)
   ....: 
10 loops, best of 3: 45.9 ms per loop

用 numpy tensordot 进行张量乘法

Tensor multiplication with numpy tensordot

python

performance

numpy

matrix

linear-algebra