Поскольку описание einsum является скудным в документации по горелке, я решил написать этот пост в документе, сравнить и сопоставить поведение torch.einsum() по сравнению с numpy.einsum().
Различия:
- NumPy допускает как строчные, так и прописные буквы
[a-zA-Z] для «строки индекса», тогда как PyTorch допускает только строчные буквы [a-z].
- NumPy поддерживает множество аргументов ключевых слов (например,
optimize) в дополнение к nd-arrays, тогда как PyTorch не обеспечивает такой гибкости
Вот реализации некоторых примеров как в PyTorch, так и в NumPy:
# input tensors
In [16]: vec
Out[16]: tensor([0, 1, 2, 3])
In [17]: aten
Out[17]:
tensor([[11, 12, 13, 14],
[21, 22, 23, 24],
[31, 32, 33, 34],
[41, 42, 43, 44]])
In [18]: bten
Out[18]:
tensor([[1, 1, 1, 1],
[2, 2, 2, 2],
[3, 3, 3, 3],
[4, 4, 4, 4]])
1) Умножение матриц
PyTorch: torch.matmul(aten, bten); aten.mm(bten)
NumPy: np.einsum("ij, jk -> ik", arr1, arr2)
In [19]: torch.einsum('ij, jk -> ik', aten, bten)
Out[19]:
tensor([[130, 130, 130, 130],
[230, 230, 230, 230],
[330, 330, 330, 330],
[430, 430, 430, 430]])
2) Извлечение элементов по главной диагонали
PyTorch: torch.diag(aten)
NumPy: np.einsum("ii -> i", arr)
In [28]: torch.einsum('ii -> i', aten)
Out[28]: tensor([11, 22, 33, 44])
3) произведение Адамара (т.е. поэлементное произведение двух тензоров)
PyTorch: aten * bten
NumPy: np.einsum("ij, ij -> ij", arr1, arr2)
In [34]: torch.einsum('ij, ij -> ij', aten, bten)
Out[34]:
tensor([[ 11, 12, 13, 14],
[ 42, 44, 46, 48],
[ 93, 96, 99, 102],
[164, 168, 172, 176]])
4) Поэлементное возведение в квадрат
PyTorch: aten ** 2
NumPy: np.einsum("ij, ij -> ij", arr, arr)
In [37]: torch.einsum('ij, ij -> ij', aten, aten)
Out[37]:
tensor([[ 121, 144, 169, 196],
[ 441, 484, 529, 576],
[ 961, 1024, 1089, 1156],
[1681, 1764, 1849, 1936]])
Общее : По элементам nth power может быть реализована путем повторения строки индекса и тензора n раз.
Например, вычисление поэлементной 4-й степени тензора может быть выполнено с использованием:
# NumPy: np.einsum('ij, ij, ij, ij -> ij', arr, arr, arr, arr)
In [38]: torch.einsum('ij, ij, ij, ij -> ij', aten, aten, aten, aten)
Out[38]:
tensor([[ 14641, 20736, 28561, 38416],
[ 194481, 234256, 279841, 331776],
[ 923521, 1048576, 1185921, 1336336],
[2825761, 3111696, 3418801, 3748096]])
5) Трасса (то есть сумма элементов главной диагонали)
PyTorch: torch.trace(aten)
NumPy einsum: np.einsum("ii -> ", arr)
In [44]: torch.einsum('ii -> ', aten)
Out[44]: tensor(110)
6) Транспонирование матрицы
PyTorch: torch.transpose(aten, 1, 0)
NumPy einsum: np.einsum("ij -> ji", arr)
In [58]: torch.einsum('ij -> ji', aten)
Out[58]:
tensor([[11, 21, 31, 41],
[12, 22, 32, 42],
[13, 23, 33, 43],
[14, 24, 34, 44]])
7) Внешний продукт (векторов)
PyTorch: torch.ger(vec, vec)
NumPy einsum: np.einsum("i, j -> ij", vec, vec)
In [73]: torch.einsum('i, j -> ij', vec, vec)
Out[73]:
tensor([[0, 0, 0, 0],
[0, 1, 2, 3],
[0, 2, 4, 6],
[0, 3, 6, 9]])
8) Внутренний продукт (векторов)
PyTorch: torch.ger(vec1, vec2)
NumPy einsum: np.einsum("i, i -> ", vec1, vec2)
In [76]: torch.einsum('i, i -> ', vec, vec)
Out[76]: tensor(14)
9) Сумма по оси 0
PyTorch: torch.sum(aten, 0)
NumPy einsum: np.einsum("ij -> j", arr)
In [85]: torch.einsum('ij -> j', aten)
Out[85]: tensor([104, 108, 112, 116])
10) Сумма по оси 1
PyTorch: torch.sum(aten, 1)
NumPy einsum: np.einsum("ij -> i", arr)
In [86]: torch.einsum('ij -> i', aten)
Out[86]: tensor([ 50, 90, 130, 170])
11) Пакетное умножение матриц
PyTorch: torch.bmm(batch_ten, batch_ten)
NumPy: np.einsum("bij, bjk -> bik", batch_ten, batch_ten)
In [90]: batch_ten = torch.stack((aten, bten), dim=0)
In [91]: batch_ten
Out[91]:
tensor([[[11, 12, 13, 14],
[21, 22, 23, 24],
[31, 32, 33, 34],
[41, 42, 43, 44]],
[[ 1, 1, 1, 1],
[ 2, 2, 2, 2],
[ 3, 3, 3, 3],
[ 4, 4, 4, 4]]])
In [92]: batch_ten.shape
Out[92]: torch.Size([2, 4, 4])
# batch matrix multiply using einsum
In [96]: torch.einsum("bij, bjk -> bik", batch_ten, batch_ten)
Out[96]:
tensor([[[1350, 1400, 1450, 1500],
[2390, 2480, 2570, 2660],
[3430, 3560, 3690, 3820],
[4470, 4640, 4810, 4980]],
[[ 10, 10, 10, 10],
[ 20, 20, 20, 20],
[ 30, 30, 30, 30],
[ 40, 40, 40, 40]]])
12) Сумма по оси 2
PyTorch: torch.sum(batch_ten, 2)
NumPy einsum: np.einsum("ijk -> ij", arr3D)
In [99]: torch.einsum("ijk -> ij", batch_ten)
Out[99]:
tensor([[ 50, 90, 130, 170],
[ 4, 8, 12, 16]])
13) Суммировать все элементы в nD-тензоре
PyTorch: torch.sum(batch_ten)
NumPy einsum: np.einsum("ijk -> ", arr3D)
In [101]: torch.einsum("ijk -> ", batch_ten)
Out[101]: tensor(480)
14) Сумма по нескольким осям (т.е. маргинализация)
PyTorch: torch.sum(arr, dim=(dim0, dim1, dim2, dim3, dim4, dim6, dim7))
NumPy: np.einsum("ijklmnop -> n", nDarr)
# 8D tensor
In [103]: nDten = torch.randn((3,5,4,6,8,2,7,9))
In [104]: nDten.shape
Out[104]: torch.Size([3, 5, 4, 6, 8, 2, 7, 9])
# marginalize out dimension 5 (i.e. "n" here)
In [111]: esum = torch.einsum("ijklmnop -> n", nDten)
In [112]: esum
Out[112]: tensor([ 98.6921, -206.0575])
# marginalize out axis 5 (i.e. sum over rest of the axes)
In [113]: tsum = torch.sum(nDten, dim=(0, 1, 2, 3, 4, 6, 7))
In [115]: torch.allclose(tsum, esum)
Out[115]: True
15) Продукты с двойными точками (аналогично torch.sum (hadamard-product) см. 3)
PyTorch: torch.sum(aten * bten)
NumPy: np.einsum("ij, ij -> ", arr1, arr2)
In [120]: torch.einsum("ij, ij -> ", aten, bten)
Out[120]: tensor(1300)