双点产品，以numpy广播

Question

我有以下操作：

import numpy as np

x = np.random.rand(3,5,5)
w = np.random.rand(5,5)

y=np.zeros((3,5,5))
for i in range(3):
    y[i] = np.dot(w.T,np.dot(x[i],w))

对应于伪表达式y[m,i,j] = sum( w[k,i] * x[m,k,l] * w[l,j], axes=[k,l]或等效只是wT ， x ， w广播商在x的第一维上的点积。

如何使用numpy的广播规则实施它？

提前致谢。

Answer 1

这是一种使用np.tensordot的矢量化方法，应该比np.tensordot broadcasting + summation都要好-

# Take care of "np.dot(x[i],w)" term
x_w = np.tensordot(x,w,axes=((2),(0)))

# Perform "np.dot(w.T,np.dot(x[i],w))" : "np.dot(w.T,x_w)"
y_out = np.tensordot(x_w,w,axes=((1),(0))).swapaxes(1,2)

另外，所有的混乱都可以通过一次np.einsum调用来解决，但可能会更慢-

y_out = np.einsum('ab,cae,eg->cbg',w,x,w)

运行时测试-

In [114]: def tensordot_app(x, w):
     ...:     x_w = np.tensordot(x,w,axes=((2),(0)))
     ...:     return np.tensordot(x_w,w,axes=((1),(0))).swapaxes(1,2)
     ...: 
     ...: def einsum_app(x, w):
     ...:     return np.einsum('ab,cae,eg->cbg',w,x,w)
     ...: 

In [115]: x = np.random.rand(30,50,50)
     ...: w = np.random.rand(50,50)
     ...: 

In [116]: %timeit tensordot_app(x, w)
1000 loops, best of 3: 477 µs per loop

In [117]: %timeit einsum_app(x, w)
1 loop, best of 3: 219 ms per loop

---

给广播一个机会

总和是-

y[m,i,j] = sum( w[k,i] * x[m,k,l] * w[l,j], axes=[k,l] )

因此，这三个术语将堆叠起来进行广播，就像这样-

w : [ N x k x i x N x N]
x : [ m x k x N x l x N]
w : [ N x N X N x l x j]

，其中N表示要添加新轴，以方便沿这些暗处进行broadcasting 。

然后将新轴添加到None/np.newaxis的术语看起来像这样-

w : w[None, :,    :,    None, None]
x : x[:,    :,    None, :,    None]
w : w[None, None, None, :,       :]

因此，广播的产品将是-

p = w[None,:,:,None,None]*x[:,:,None,:,None]*w[None,None,None,:,:]

最后，输出将是减总和以损失(k,l) ，即轴= (1,3) -

y = p.sum((1,3))