使用 numpy 进行批量张量乘法

Question

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

我有一个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，但到目前为止没有成功。我错过了什么？

Answer 1

由于缓存的使用在小张量序列上不是问题（就像对于大型矩阵的一般点积一样），所以很容易用简单的循环来表述问题。

示例

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

Answer 2

我们可以使用tensor matrix-multiplication 与np.tensordot 和einsum 的组合，基本上分两步完成-

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

基准测试

基于 OP 的 cmets 设置：

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

Answer 3

您可以或多或少直接将您的公式转换为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.]])

Answer 4

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

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

让X 成为您想要批量处理的数据的m x d 矩阵，并让Z 成为您想要的m x d 结果。我们将到达Z.T（转置），因为它更容易。

请注意，为了得到R 贡献的方程式，我们可以写成

然后我们可以将其作为一个 numpy 矩阵乘以 R.dot(X.T)。

同样，观察T 的贡献是

括号内是T 和X.T 之间的批处理矩阵乘法。因此，如果我们将括号内的数量定义为

我们可以在 numpy 中使用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 时的 einsum 函数快大约 3-4 倍。如果我们要避免使用尽可能多的转置，它可能会更快。