python之使用 numpy 的批量张量乘法

我正在尝试执行以下矩阵和张量乘法，但是是批处理的。

我有一个 x 向量列表:

x = np.array([[2.0, 2.0], [3.0, 3.0], [4.0, 4.0], [5.0, 5.0]]) 

以及以下矩阵和张量:

R = np.array( 
    [ 
        [1.0, 1.0], 
        [0.0, 1.0], 
    ] 
) 
T = np.array( 
    [ 
        [ 
            [2.0, 0.0], 
            [0.0, 0.0], 
        ], 
        [ 
            [0.0, 0.0], 
            [0.0, 2.0], 
        ] 
    ] 
) 

批量矩阵乘法相对简单:

x.dot(R.T) 

然而，我正在为第二部分而苦苦挣扎。

我尝试使用 tensordot 但到目前为止没有成功。我错过了什么？

请您参考如下方法:

由于缓存使用不是小张量序列的问题(就像大矩阵的一般点积一样)，因此很容易用简单的循环来表述问题。

示例

import numba as nb 
import numpy as np 
import time 
 
@nb.njit(fastmath=True,parallel=True) 
def tensor_mult(T,x): 
  res=np.empty((x.shape[0],T.shape[0]),dtype=T.dtype) 
  for l in nb.prange(x.shape[0]): 
    for i in range(T.shape[0]): 
      sum=0. 
      for j in range(T.shape[1]): 
        for k in range(T.shape[2]): 
          sum+=T[i,j,k]*x[l,j]*x[l,k] 
      res[l,i]=sum 
  return res 

基准测试

x = np.random.rand(1000000,6) 
T = np.random.rand(6,6,6) 
 
#first call has a compilation overhead (about 0.6s) 
res=tensor_mult(T,x) 
 
t1=time.time() 
for i in range(10): 
  #@divakar 
  #Tx = np.tensordot(T,x,axes=((1),(1))) 
  #out = np.einsum('ikl,lk->li',Tx,x) 
 
  res=tensor_mult(T,x) 
 
print(time.time()-t1) 

成绩(4C/8T)

Divakars solution: 191ms 
Simple loops: 62.4ms