在 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 broadcaster 在 [=13 的第一个维度上的简单点积=].

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

提前致谢。

Answer 1

这是一种矢量化方法 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))

在 numpy 中广播的双点积

Double dot product with broadcasting in numpy

python

numpy

matrix

tensor