带numpy的成批张量乘法

由于高速缓存的使用不是一个小张量序列的问题(对于大矩阵的一般点积也是如此),所以用简单的循环来描述这个问题是很容易的。

例子

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

我们可以结合使用 tensor matrix-multiplication 具有 np.tensordot 和 einsum 基本上分两步做-

Tx = np.tensordot(T,x,axes=((1),(1)))
out = np.einsum('ikl,lk->li',Tx,x)

标杆管理

根据OP的意见设置:

In [1]: import numpy as np

In [2]: x = np.random.rand(1000000,6)

In [3]: T = np.random.rand(6,6,6)

时间安排-

# @Han Altae-Tran's soln
In [4]: %%timeit
   ...: W = np.matmul(T,x.T) 
   ...: ZT = np.sum(W*x.T[np.newaxis,:,:], axis=1).T
   ...: 
1 loops, best of 3: 496 ms per loop

# @Paul Panzer's soln-1
In [5]: %timeit np.einsum('ijk,lj,lk->li', T, x, x)
1 loops, best of 3: 831 ms per loop

# @Paul Panzer's soln-2
In [6]: %timeit ((x[:, None, None, :]@T).squeeze()@x[..., None]).squeeze()
1 loops, best of 3: 1.39 s per loop

# @Paul Panzer's soln-3
In [7]: %timeit np.einsum('ijl,lj->li', T@x.T, x)
1 loops, best of 3: 358 ms per loop

# From this post's soln
In [8]: %%timeit
   ...: Tx = np.tensordot(T,x,axes=((1),(1)))
   ...: out = np.einsum('ikl,lk->li',Tx,x)
   ...: 
1 loops, best of 3: 168 ms per loop

您可以或多或少地直接将公式转换为 einsum 以下内容:

>>> np.einsum('ijk,lj,lk->li', T, x, x)
array([[ 8.,  8.],
       [18., 18.],
       [32., 32.],
       [50., 50.]])

只使用 @ 以下内容:

>>> ((x[:, None, None, :]@T).squeeze()@x[..., None]).squeeze()
array([[ 8.,  8.],
       [18., 18.],
       [32., 32.],
       [50., 50.]])

或者混合动力车:

>>> np.einsum('ijl,lj->li', T@x.T, x)
array([[ 8.,  8.],
       [18., 18.],
       [32., 32.],
       [50., 50.]])

正如paul所指出的,einsum是完成任务的一个简单方法,但是如果速度是一个问题,那么通常最好坚持使用典型的numpy函数。

这可以通过写出方程式并将步骤转换为矩阵运算来实现。

让 X 成为 m x d 要批处理的数据矩阵,以及 Z 成为 M x D码 你想要的结果。我们将到达 Z.T (转置)因为它更容易。

注意,为了得到 R 我们可以写的贡献

然后我们可以把它当作一个numpy矩阵乘法 R.dot(X.T) 是的。

同样,请注意 T 贡献是

括号内是一个批处理矩阵 T型 和 X.T 是的。因此,如果我们将括号内的数量定义为

我们可以在纽比用 W = np.matmul(T,X.T) 是的。继续我们的简化,我们看到 T型 贡献是

相当于 np.sum(W*X.T[np.newaxis,:,:], axis=1) 是的。把所有的东西放在一起,我们最终

W = np.matmul(T,X.T) 
ZT = R.dot(X.T) + np.sum(W*X.T[np.newaxis,:,:], axis=1) 
Z = ZT.T

对于较大的批处理大小,当 d=2 是的。如果我们要避免使用那么多转置,可能会更快一些。