Пакетное тензорное умножение с numy

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

Поскольку использование кэша не является проблемой для последовательности малых тензоров (как это было бы для общих точечных произведений больших матриц), проблему легко сформулировать с помощью простых циклов.

* ** 1003 тысяча два * Пример

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)

Бенчмаркинг

Настройка на основе комментариев ОП:

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 - это простой способ выполнить задачу, но если скорость вызывает беспокойство, то, как правило, лучше придерживаться типичных функций numpy.

Это может быть достигнуто путем выписывания уравнения и перевода шагов в матричные операции.

Пусть X будет m x d матрицей данных, для которой вы хотите выполнить пакетную обработку, а Z будет m x d желаемым результатом. Мы придем к Z.T (транспонировать), потому что это проще.

Обратите внимание, что для получения уравнения для вклада R мы можем написать

Тогда мы можем это умножить на числовую матрицу 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

Для больших размеров партии это примерно в 3-4 раза быстрее, чем функция einsum при d=2. Если бы нам не пришлось использовать как можно больше транспозов, это могло бы быть даже немного быстрее.

Answer 4

Вы можете более или менее напрямую перевести свою формулу в 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.]])