在 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] 或等效于w.T、x、w 广播公司在x 的第一个维度上的点积。

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

提前致谢。

Answer 1

这是np.tensordot 的一种矢量化方法，它应该比broadcasting + summation anyday -

# 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,:,:]

最后，输出将是 sum-reduction to loss (k,l)，即轴 =(1,3) -

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