numpy广播的双点产品

时间:2017-08-16 16:24:15

标签: python numpy matrix tensor

我有以下操作:

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.Txw广播员在第x的第一维上的点积。

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

提前致谢。

1 个答案:

答案 0 :(得分:2)

这是一个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))