Python:使用Tensordot进行张量x矩阵乘法

时间:2011-03-04 18:51:57

标签: python numpy

您希望以下列方式将张量乘以矩阵:

尺寸

W:a x b x c

V:a x c

我希望Z这样

Z[i]=dot(W[i],V[i])

Z的维度为a x ( (b x c) . (c x 1)),因此(a x b)

我已尝试numpy.tensordot这样做,但还是没能。它可以做我想要的吗?如果不是,我怎么能这样做没有循环。

基本上相当于

def f(W,V):    
    Z=[]    
    for i in range(len(W)):    
        Z.append(dot(W[i],V[i]))    
    return Z

由于

编辑:具体是这可以通过tensordot实现吗?

2 个答案:

答案 0 :(得分:5)

np.einsum("abc,ac -> ab", w, v)

import numpy as np

def z_loop(w,v): # define it to check that `einsum()` gives necessary result
    z = np.empty(w.shape[:-1], dtype=w.dtype)
    for i in range(z.shape[0]):
        z[i,:] = np.dot(w[i,:], v[i,:])
    return z

w = np.random.uniform(size=(3,4,5))
v = np.random.uniform(size=w.shape[::2])
assert np.allclose(z_loop(w, v), np.einsum('abc,ac -> ab', w, v))

可能有更简单的变体(通过dot().reshape()),但einsum()对于任务描述最为明显。

def z_dot(w, v):
    z = np.dot(w, v[:,...,np.newaxis])
    z = z.reshape(z.shape[:-1])
    return np.diagonal(z, axis2=-1).T

assert np.allclose(z_dot(w, v), np.einsum('abc,ac -> ab', w, v))

答案 1 :(得分:2)

怎么样?
import numpy as np

a,b,c=3,5,6

r=np.random.random
W = r((a,b,c))
V = r((a,c))

Z = np.sum(W*V[:,np.newaxis,:],axis=2)

不使用循环或更新的功能,并且应该相当快。与J.F. Sebastian的帖子中的“z_loop”相比:

print np.sum(np.abs(Z-z_loop(W,V)))

给出4.99600361081e-16