问题描述
我有以下形状的 numpy 数组:(2000,3) (3,2000,2)。我想知道以下张量收缩是否:
t1 = np.einsum("ij,jik->ik",a,b)
可以使用np.dot
来完成。我试过这样做:
out = np.zeros((2000,2))
for k in range(2):
out[:,k] = np.dot(a,b[:,:,k])
但它抛出一个错误: ValueError: 无法将输入数组从形状 (1000,1000) 广播到形状 (1000)
迭代第一个维度 (i) 似乎效率不高。
解决方法
举一个最小的例子(我们喜欢在问题中看到):
In [173]: a = np.arange(12).reshape(4,3); b=np.arange(24).reshape(3,4,2)
In [174]: t1 = np.einsum('ij,jik->ik',a,b)
In [175]: t1.shape
Out[175]: (4,2)
In [176]: t1
Out[176]:
array([[ 40,43],[136,148],[268,289],[436,466]])
工作循环,i
上的迭代(是的,即使它可能很大)
In [177]: out = np.zeros((4,2),int)
In [178]: for i in range(4):
...: out[i,:] = np.dot(a[i],b[:,i,:])
...:
In [179]: out
Out[179]:
array([[ 40,466]])
在我们担心效率之前,它必须工作!
现在,如果我们转置 b
,我们可以将 einsum
更改为:
In [181]: np.einsum('ij,ijk->ik',b.transpose(1,2)).shape
Out[181]: (4,2)
在正确的位置使用 j
,dot
可以工作 - 但它在共享的 outer
维度上执行 i
产品。我们必须取对角线:
In [182]: np.dot( a,2)).shape
Out[182]: (4,2)
添加 matmul/@
以将此“外部”行为更改为更常见的“批处理”。但是我们需要为 a
添加一个维度,因此 i
是所有项的 3 个中的第一个:
In [184]: np.matmul( a[:,None,:],2)).shape
Out[184]: (4,1,2)
最后挤出1维:
In [186]: np.matmul( a[:,2))[:,:]
Out[186]:
array([[ 40,466]])
最后一个表达式应该和 einsum
一样快,如果不是更快的话。
In [192]: timeit t1 = np.einsum('ij,b)
8.05 µs ± 63.4 ns per loop (mean ± std. dev. of 7 runs,100000 loops each)
In [193]: timeit np.matmul( a[:,:]
5.19 µs ± 183 ns per loop (mean ± std. dev. of 7 runs,100000 loops each)