torch.einsum

torch.einsum(equation, *operands) → Tensor [source]

Вычисляет сумму произведений элементов входных operands по измерениям, указанным с помощью обозначений, основанных на соглашении об обозначении Эйнштейна.

Einsum позволяет вычислять многие общие многомерные линейно-алгебраические операции над массивами, представляя их в краткой форме, основанной на соглашении об обозначении Эйнштейна, заданном equation. Подробное описание этого формата приведено ниже, но общая идея заключается в маркировке каждого измерения входного operands некоторой подстрокой и определении того, какие подстроки являются частью выходных данных. Затем выходные данные вычисляются путём суммирования произведения элементов operands по измерениям, подстроки которых не являются частью выходных данных. Например, умножение матриц может быть вычислено с помощью einsum, как torch.einsum(“ij,jk->ik”, A, B). Здесь j — подстрока суммирования, а i и k — подстроки выходных данных (подробнее о том, почему, см. раздел ниже).

Уравнение:

Строка equation определяет подстроки (буквы в [a-zA-Z]) для каждого измерения входного operands в том же порядке, что и измерения, отделяя подстроки для каждого операнда запятой (','), например, ‘ij,jk’ задаёт подстроки для двух 2D-операндов. Измерения, помеченные одной и той же подстрокой, должны быть совместимы с трансляцией, то есть их размер должен либо совпадать, либо быть 1. Исключением является случай, когда подстрока повторяется для одного и того же операнда входных данных, в этом случае размеры измерений, помеченных этой подстрокой для этого операнда, должны совпадать, и операнд будет заменён его диагональю по этим измерениям. Подстроки, которые появляются ровно один раз в equation , будут частью выходных данных, отсортированных в порядке возрастания алфавитного порядка. Выходные данные вычисляются путём поэлементного умножения входных operands с выравниванием их измерений на основе подстрок, а затем суммированием измерений, подстроки которых не являются частью выходных данных.

Необязательно, подстроки выходных данных могут быть явно определены путём добавления стрелки ('->') в конце уравнения, за которой следуют подстроки для выходных данных. Например, следующее уравнение вычисляет транспонирование умножения матриц: 'ij,jk->ki'. Подстроки выходных данных должны появляться по меньшей мере один раз для некоторого операнда входных данных и не более одного раза для выходных данных.

Эллипсис ('...') может быть использован вместо подстрок для трансляции измерений, покрытых эллипсисом. Каждый операнд входных данных может содержать не более одного эллипса, который будет покрывать измерения, не покрытые подстроками, например, для операнда входных данных с 5 измерениями, эллипсы в уравнении ‘ab…c’ покрывают третье и четвёртое измерения. Эллипсис не обязан покрывать одинаковое количество измерений в operands, но «форма» эллипса (размер измерений, покрытых ими) должна быть совместима с трансляцией. Если выходные данные не определены явно с помощью обозначения стрелки ('->'), эллипсис будет стоять первым в выходных данных (левые измерения), перед метками подстрок, которые появляются ровно один раз для операндов входных данных. Например, следующее уравнение реализует умножение матриц в пакете ‘…ij,…jk’.

Несколько заключительных замечаний: уравнение может содержать пробелы между разными элементами (подстроки, эллипсы, стрелка и запятая), но что-то вроде ‘…’ не является допустимым. Пустая строка ‘’ допустима для скалярных операндов.

Примечание

torch.einsum обрабатывает эллипсы ('...') по-другому, чем NumPy, так как допускает суммирование измерений, покрытых эллипсом, то есть эллипсы не обязательно должны быть частью выходных данных.

Примечание

Эта функция использует opt_einsum (https://optimized-einsum.readthedocs.io/en/stable/) для ускорения вычислений или снижения потребления памяти путём оптимизации порядка сокращения. Эта оптимизация происходит, когда есть по крайней мере три входных данных, поскольку порядок не имеет значения в противном случае. Обратите внимание, что нахождение _оптимального_ пути является задачей NP-трудной, поэтому opt_einsum полагается на различные эвристики для достижения почти оптимальных результатов. Если opt_einsum недоступен, по умолчанию порядок сокращения — слева направо.

Чтобы обойти это поведение по умолчанию, добавьте следующую строку, чтобы отключить использование opt_einsum и пропустить вычисление пути: torch.backends.opt_einsum.enabled = False

Чтобы указать, какой стратегии вы хотите, чтобы opt_einsum вычислял путь сокращения, добавьте следующую строку: torch.backends.opt_einsum.strategy = ‘auto’. По умолчанию используется стратегия «auto», и мы также поддерживаем «greedy» и «optimal». Обратите внимание, что время работы «optimal» факториально зависит от количества входных данных! Более подробную информацию см. в документации opt_einsum (https://optimized-einsum.readthedocs.io/en/stable/path_finding.html).

Примечание

Начиная с PyTorch 1.10 torch.einsum() также поддерживает формат списка подсписков (см. примеры ниже). В этом формате подстроки для каждого операнда указываются подсписками, списком целых чисел в диапазоне [0, 52). Эти подсписки следуют за их операндами, и дополнительный подсписок может появиться в конце входных данных, чтобы указать подстроки выходных данных, например, torch.einsum(op1, sublist1, op2, sublist2, …, [subslist_out]). Объект Python Ellipsis может быть предоставлен в подсписке, чтобы включить трансляцию, как описано в разделе «Уравнение» выше.

Параметры

• уравнение (str) – Подстроки для вычисления суммирования Эйнштейна.
• операнды (Список[Tensor]) – Те тензоры, для которых требуется вычислить суммирование Эйнштейна.

Тип возвращаемого значения

Tensor

Примеры:

>>> # trace
>>> torch.einsum('ii', torch.randn(4, 4))
tensor(-1.2104)

>>> # diagonal
>>> torch.einsum('ii->i', torch.randn(4, 4))
tensor([-0.1034,  0.7952, -0.2433,  0.4545])

>>> # outer product
>>> x = torch.randn(5)
>>> y = torch.randn(4)
>>> torch.einsum('i,j->ij', x, y)
tensor([[ 0.1156, -0.2897, -0.3918,  0.4963],
        [-0.3744,  0.9381,  1.2685, -1.6070],
        [ 0.7208, -1.8058, -2.4419,  3.0936],
        [ 0.1713, -0.4291, -0.5802,  0.7350],
        [ 0.5704, -1.4290, -1.9323,  2.4480]])

>>> # batch matrix multiplication
>>> As = torch.randn(3, 2, 5)
>>> Bs = torch.randn(3, 5, 4)
>>> torch.einsum('bij,bjk->bik', As, Bs)
tensor([[[-1.0564, -1.5904,  3.2023,  3.1271],
        [-1.6706, -0.8097, -0.8025, -2.1183]],

        [[ 4.2239,  0.3107, -0.5756, -0.2354],
        [-1.4558, -0.3460,  1.5087, -0.8530]],

        [[ 2.8153,  1.8787, -4.3839, -1.2112],
        [ 0.3728, -2.1131,  0.0921,  0.8305]]])

>>> # with sublist format and ellipsis
>>> torch.einsum(As, [..., 0, 1], Bs, [..., 1, 2], [..., 0, 2])
tensor([[[-1.0564, -1.5904,  3.2023,  3.1271],
        [-1.6706, -0.8097, -0.8025, -2.1183]],

        [[ 4.2239,  0.3107, -0.5756, -0.2354],
        [-1.4558, -0.3460,  1.5087, -0.8530]],

        [[ 2.8153,  1.8787, -4.3839, -1.2112],
        [ 0.3728, -2.1131,  0.0921,  0.8305]]])

>>> # batch permute
>>> A = torch.randn(2, 3, 4, 5)
>>> torch.einsum('...ij->...ji', A).shape
torch.Size([2, 3, 5, 4])

>>> # equivalent to torch.nn.functional.bilinear
>>> A = torch.randn(3, 5, 4)
>>> l = torch.randn(2, 5)
>>> r = torch.randn(2, 4)
>>> torch.einsum('bn,anm,bm->ba', l, A, r)
tensor([[-0.3430, -5.2405,  0.4494],
        [ 0.3311,  5.5201, -3.0356]])