numpy的批量张量乘法

时间:2018-06-26 15:49:37

标签: python numpy linear-algebra

我正在尝试执行以下矩阵和张量乘法,但是是分批进行的。

matrix and tensor multiplication

我有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,但到目前为止没有成功。我想念什么?

4 个答案:

答案 0 :(得分:2)

我们可以将tensor matrix-multiplicationnp.tensordoteinsum结合使用,基本上可以分两步进行-

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

答案 1 :(得分:1)

您可以或多或少直接将公式转换为cols <- lapply(all_lists, function(l) c(l, rep(NA, len - length(l)))) as.data.frame(Reduce(cbind, cols, init = NULL)) #> V1 V2 V3 V4 V5 #> 1 USA Microsoft NA 1989 CEO #> 2 Singapore University of London NA 2001 Chairman #> 3 UK Boeing NA 2018 VP of sales #> 4 NA Apple NA NA General Manager #> 5 NA NA NA NA Director

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

答案 2 :(得分:1)

如Paul所指出的那样,einsum是完成任务的一种简便方法,但是如果考虑到速度,那么通常最好还是使用典型的numpy函数。

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

X是要分批处理的m x d数据矩阵,而Z是您想要的m x d结果。我们会更容易到达Z.T(移调)。

请注意,为了得出R贡献的等式,我们可以编写

enter image description here

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

类似地,观察到T的贡献是

enter image description here

括号内是TX.T之间乘以的批处理矩阵。因此,如果我们定义括号内的数量为

enter image description here

我们可以使用W = np.matmul(T,X.T)以numpy的形式实现。继续简化,我们看到T的贡献是

enter image description here

相当于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倍。如果我们要避免使用尽可能多的转置,它可能会更快一些。

答案 3 :(得分: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