torch.einsum详解

torch.einsum 是 PyTorch 中用于执行高效张量运算的函数，基于爱因斯坦求和约定（Einstein summation convention）。它能够处理复杂的张量操作，并简化代码书写。

基本语法

torch.einsum(subscripts, *operands)

• subscripts：一个字符串，用于描述输入张量的维度如何结合。
• *operands：待操作的张量。

爱因斯坦求和约定

爱因斯坦求和约定是一个简化张量运算的方式，省略了显式的求和符号。通过指定各维度的标签，可以直接描述复杂的张量运算。

语法结构

• "nqhd,nkhd->nhqk": 这个字符串描述了如何对两个张量进行操作，并生成输出张量的维度。

  • n：批次大小（batch size）
  • q：查询序列长度（query length）
  • k：键序列长度（key length）
  • h：注意力头的数量（number of heads）
  • d：每个注意力头的维度（dimension per head）

示例代码

以下是使用 torch.einsum 计算多头注意力机制中点积相似性的示例代码：

import torch

# 定义多头注意力机制的点积计算函数
def compute_attention_scores(queries, keys):
    # 计算点积相似性分数
    energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
    return energy

# 示例数据
N = 1            # 批次大小
q =</