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

0 投票
2 回答
2276 浏览
提问于 2025-04-16 13:01

你好,我正在尝试以以下方式将一个张量和一个矩阵相乘:

维度

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 个回答

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

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))

撰写回答