张量乘法与numpy tensordot

Question

我有张量U由n维矩阵（d，k）和维数（k，n）的矩阵V组成。

我想将它们相乘，以便结果返回一个维度矩阵（d，n），其中列j是U的矩阵j和V的列j之间的矩阵乘法的结果。

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

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对于执行复杂的张量收缩非常有用，但它可能需要一段时间才能完全包围它的工作方式。您应该查看文档中的示例（上面链接）。

其他一些选择：

1. 使用broadcasting进行元素乘法，然后求和：

   res3 = (U * V[None, ...]).sum(1)
   

2. 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