张量乘法与numpy tensordot

Question

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

我想将它们相乘,以便结果返回一个维度矩阵(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使用爱因斯坦符号表示张量收缩.在该表达式中'ijk,jk->ik'的上方,i,j和k是对应于不同尺寸的下标U和V.每个以逗号分隔的分组对应于传递给的操作数之一np.einsum (在这种情况下U具有维度ijk并V具有维度jk).该'->ik'部件指定输出数组的尺寸.将输出字符串中不存在的具有下标的任何维度相加.

np.einsum对于执行复杂的张量收缩非常有用,但它可能需要一段时间才能完全包围它的工作方式.您应该查看文档中的示例(上面链接).

其他一些选择:

1. 元素乘法与广播,然后求和:

   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