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
字符串指定输入operands
每个维度的下标([a-zA-Z] 中的字母),顺序与维度相同,并用逗号 (',') 分隔每个操作数的下标,例如 ‘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-hard 问题,因此,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 对象可以提供给子列表以启用广播,如上文“方程”部分所述。示例
>>> # 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]])