torch.onnx.ops#

创建于：2025年6月10日 | 最后更新于：2025年6月20日

将ONNX运算符作为原生 torch.fx 运算符。

此模块提供了一组函数，用于在 FX 图中创建可导出到 ONNX 的 ONNX 运算符。

符号运算符#

可用于在 FX 图中以符号方式创建任何 ONNX 运算符的运算符。这些运算符不执行实际计算。建议您在 if torch.onnx.is_in_onnx_export 块内使用它们。

torch.onnx.ops.symbolic(domain_op, /, inputs, attrs=None, *, dtype, shape, version=None, metadata_props=None)[source]#

创建一个符号 FX 运算符来表示任意 ONNX 运算符。

此函数用于创建具有单个输出的符号运算符。要创建具有多个输出的运算符，请使用 symbolic_multi_out()。

您可以 if torch.onnx.is_in_onnx_export() 来仅在 torch.onnx.export() 期间有条件地启用符号逻辑。

示例

class CustomOp(torch.nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Normal torch operators can interleave with the symbolic ops during ONNX export
        x = x + 1

        # Create a symbolic ONNX operator with the name "CustomOp" in the "custom_domain" domain.
        # The output tensor will have the specified dtype and shape
        val = torch.onnx.ops.symbolic(
            "custom_domain::CustomOp",
            (x,),
            dict(attr_key="attr_value"),
            dtype=x.dtype,
            shape=x.shape,
            version=1,
        )

        # The result of the symbolic op can be used in normal torch operations during ONNX export
        return torch.nn.functional.relu(val)


# You may then export this model to ONNX using torch.onnx.export(..., dynamo=True).

参数

domain_op (str) – 域和运算符名称，用“::”分隔。例如，“custom_domain::CustomOp”。
inputs (Sequence[torch.Tensor | None]) – 运算符的输入张量。
attrs (dict[str, int | float | str | bool | Sequence[int] | Sequence[float] | Sequence[str] | Sequence[bool]] | None) – 运算符的属性。键是属性名称，值是属性值。有效的属性类型是 int、float、str、bool 以及 int、float、str 和 bool 的列表。不支持张量属性。
dtype (torch.dtype | int) – 输出张量的数据类型。可以是 torch.dtype 或代表 ONNX 数据类型的整数。
shape (Sequence[int | torch.SymInt]) – 输出张量的形状。这可以是一个整数或 SymInt 值的列表。
version (int | None) – 用于该运算符的 opset 版本。
metadata_props (dict[str, str] | None) – ONNX 节点的元数据属性。这是一个 str-str 对的字典。

返回

运算符的输出张量。

返回类型

torch.Tensor

torch.onnx.ops.symbolic_multi_out(domain_op, /, inputs, attrs=None, *, dtypes, shapes, version=None, metadata_props=None)[source]#

创建一个符号 FX 运算符来表示具有多个输出的任意 ONNX 运算符。

您可以 if torch.onnx.is_in_onnx_export() 来仅在 torch.onnx.export() 期间有条件地启用符号逻辑。

示例

class CustomOp(torch.nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Normal torch operators can interleave with the symbolic ops during ONNX export
        x = x + 1

        # Create a symbolic ONNX operator with the name "CustomOp" in the "custom_domain" domain.
        # The output tensors will have the specified dtypes and shapes
        (out1, out2) = torch.onnx.ops.symbolic_multi_out(
            "custom_domain::CustomOp",
            (x,),
            dict(attr_key="attr_value"),
            dtypes=(x.dtype, torch.float32),
            shapes=(x.shape, [1, 2, 3]),
            version=1,
        )

        # The result of the symbolic op can be used in normal torch operations during ONNX export
        return torch.nn.functional.relu(out1 + out2)


# You may then export this model to ONNX using torch.onnx.export(..., dynamo=True).

参数

domain_op (str) – 域和运算符名称，用“::”分隔。例如，“custom_domain::CustomOp”。
inputs (Sequence[torch.Tensor | None]) – 运算符的输入张量。
attrs (dict[str, int | float | str | bool | Sequence[int] | Sequence[float] | Sequence[str] | Sequence[bool]] | None) – 运算符的属性。键是属性名称，值是属性值。有效的属性类型是 int、float、str、bool 以及 int、float、str 和 bool 的列表。不支持张量属性。
dtypes (Sequence[torch.dtype | int]) – 输出张量的数据类型。这可以是 torch.dtype 或代表 ONNX 数据类型的整数列表。此列表的长度必须等于输出的数量。
shapes (Sequence[Sequence[int | torch.SymInt]]) – 输出张量的形状。这可以是整数或 SymInt 值列表的列表。此列表的长度必须等于输出的数量。
version (int | None) – 用于该运算符的 opset 版本。
metadata_props (dict[str, str] | None) – ONNX 节点的元数据属性。这是一个 str-str 对的字典。

返回

运算符的输出张量列表。

返回类型

Sequence[torch.Tensor]

ONNX 运算符#

以下运算符已实现为原生 PyTorch 运算符，可以导出为 ONNX 运算符。它们可以在 nn.Module 中原生使用。

例如，您可以定义一个模块

class Model(torch.nn.Module):
    def forward(
        self, input_data, cos_cache_data, sin_cache_data, position_ids_data
    ):
        return torch.onnx.ops.rotary_embedding(
            input_data,
            cos_cache_data,
            sin_cache_data,
            position_ids_data,
        )

并使用以下方式将其导出到 ONNX：

input_data = torch.rand(2, 3, 4, 8)
position_ids_data = torch.randint(0, 50, (2, 3)).long()
sin_cache_data = torch.rand(50, 4)
cos_cache_data = torch.rand(50, 4)
dynamic_shapes = {
    "input_data": {0: torch.export.Dim.DYNAMIC},
    "cos_cache_data": None,
    "sin_cache_data": None,
    "position_ids_data": {0: torch.export.Dim.DYNAMIC},
}
onnx_program = torch.onnx.export(
    model,
    (input_data, cos_cache_data, sin_cache_data, position_ids_data),
    dynamic_shapes=dynamic_shapes,
    dynamo=True,
    opset_version=23,
)

打印 ONNX 程序将显示图中使用的 ONNX 运算符

<...>

graph(
    name=main_graph,
    inputs=(
        %"input_data"<FLOAT,[s0,3,4,8]>,
        %"cos_cache_data"<FLOAT,[50,4]>,
        %"sin_cache_data"<FLOAT,[50,4]>,
        %"position_ids_data"<INT64,[s0,3]>
    ),
    outputs=(
        %"rotary_embedding"<FLOAT,[s0,3,4,8]>
    ),
) {
    0 |  # rotary_embedding
        %"rotary_embedding"<FLOAT,[s0,3,4,8]> ⬅️ ::RotaryEmbedding(%"input_data", %"cos_cache_data", %"sin_cache_data", %"position_ids_data")
    return %"rotary_embedding"<FLOAT,[s0,3,4,8]>
}

以及相应的 ExportedProgram

ExportedProgram

class GraphModule(torch.nn.Module):
    def forward(self, input_data: "f32[s0, 3, 4, 8]", cos_cache_data: "f32[50, 4]", sin_cache_data: "f32[50, 4]", position_ids_data: "i64[s0, 3]"):
        rotary_embedding: "f32[s0, 3, 4, 8]" = torch.ops.onnx.RotaryEmbedding.opset23(input_data, cos_cache_data, sin_cache_data, position_ids_data);  input_data = cos_cache_data = sin_cache_data = position_ids_data = None
        return (rotary_embedding,)

torch.onnx.ops.rotary_embedding(X, cos_cache, sin_cache, position_ids=None, *, interleaved=False, num_heads=0, rotary_embedding_dim=0)[source]#

ONNX 中的 RotaryEmbedding 运算符。

https://onnx.org.cn/onnx/operators/onnx__RotaryEmbedding.html

RotaryEmbedding 是基于论文 https://arxiv.org/pdf/2104.09864 实现的旋转位置嵌入 (RoPE)。RoPE 的主要优点是它允许模型理解 token 的绝对位置以及 token 之间的相对距离。这是通过一种旋转机制实现的，其中旋转的程度根据 token 的绝对位置 (position_ids) 计算得出。在论文的旋转机制中，使用正弦和余弦函数来表示旋转角度。对于序列中的每个 token，其位置嵌入通过旋转其嵌入向量来计算。这通过将嵌入向量分成两半或交替排列每个 token 来完成，并将旋转矩阵应用于嵌入向量的每一半。旋转矩阵由 token 在序列中的位置参数化。旋转后的嵌入向量被连接起来形成每个 token 的最终位置嵌入。旋转后的位置嵌入用于自注意力机制。旋转确保模型捕获绝对和相对位置信息。

旋转机制由正弦和余弦函数定义，用于表示旋转角度。对于序列中的每个 token，其位置嵌入是通过旋转其嵌入向量来计算的。这是通过将嵌入向量分成两半或交替排列每个 token 来实现的，并将旋转矩阵应用于嵌入向量的每一半。旋转矩阵由 token 在序列中的位置参数化。旋转后的嵌入向量被连接起来形成每个 token 的最终位置嵌入。旋转后的位置嵌入用于自注意力机制。旋转确保模型捕获绝对和相对位置信息。

参数

X (Tensor) – 表示 token 嵌入的输入张量。形状为 (batch_size, num_heads, sequence_length, head_size) 的 4D 张量，或形状为 (batch_size, sequence_length, hidden_size) 的 3D 张量。对于 4D 输入张量，head_size 必须是偶数。对于 3D 输入张量，必须提供 num_heads 属性，并且 hidden_size 必须是 num_heads 的偶数倍，其中 hidden_size = num_heads * head_size
cos_cache (Tensor) – 旋转的余弦值。形状为 (max_position_id_plus_1, head_size / 2) 的 2D 张量（用于完全旋转），或形状为 (max_position_id_plus_1, rotary_embedding_dim / 2) 的 2D 张量（用于部分旋转，当提供 position_ids 时）。形状为 (batch_size, sequence_length, head_size / 2) 的 3D 张量（用于完全旋转），或形状为 (batch_size, sequence_length, rotary_embedding_dim / 2) 的 3D 张量（用于部分旋转，当不提供 position_ids 时）。max_position_id_plus_1 是模型的参数。
sin_cache (Tensor) – 旋转的正弦值。形状为 (max_position_id_plus_1, head_size / 2) 的 2D 张量（用于完全旋转），或形状为 (max_position_id_plus_1, rotary_embedding_dim / 2) 的 2D 张量（用于部分旋转，当提供 position_ids 时）。形状为 (batch_size, sequence_length, head_size / 2) 的 3D 张量（用于完全旋转），或形状为 (batch_size, sequence_length, rotary_embedding_dim / 2) 的 3D 张量（用于部分旋转，当不提供 position_ids 时）。max_position_id_plus_1 是模型的参数。
position_ids (Optional[Tensor]) – token 的位置索引。形状为 (batch_size, sequence_length) 的 2D 张量。
interleaved (bool) – 使用交错模式进行旋转。默认值为 0 (False)。
num_heads (int) – 注意力头的数量。当输入是 3D 张量时必须提供。
rotary_embedding_dim (int) – 用于应用部分旋转嵌入的旋转嵌入维度。

返回

形状与输入相同的张量。

返回类型

张量

torch.onnx.ops.attention(Q, K, V, attn_mask=None, past_key=None, past_value=None, *, is_causal=False, kv_num_heads=0, q_num_heads=0, qk_matmul_output_mode=0, scale=None, softcap=0.0, softmax_precision=None)[source]#

ONNX 中的 Attention 运算符。

https://onnx.org.cn/onnx/operators/onnx__Attention.html

在查询、键和值张量上计算缩放点积注意力，如果提供了可选的注意力掩码。

此运算符根据 K、Q 和 V 的序列长度涵盖注意力的自注意力和交叉注意力变体。

对于自注意力，kv_sequence_length 等于 q_sequence_length。

对于交叉注意力，查询和键的长度可能不同。

此运算符还涵盖以下 3 种变体，具体取决于头的数量：

多头注意力 (MHA)：如论文 https://arxiv.org/pdf/1706.03762 所述，q_num_heads = kv_num_heads。
组查询注意力 (GQA)：如论文 https://arxiv.org/pdf/2305.13245 所述，q_num_heads > kv_num_heads，q_num_heads % kv_num_heads == 0。
多查询注意力 (MQA)：如论文 https://arxiv.org/pdf/1911.02150 所述，q_num_heads > kv_num_heads，kv_num_heads=1。

基于 attn_mask 输入和 is_causal 属性计算要添加的注意力偏差，只能提供其中一个。

如果 is_causal 设置为 1，则当掩码为方形矩阵时，注意力掩码是下三角矩阵。由于对齐，注意力掩码的形式是左上角的因果偏置。
attn_mask：一个布尔掩码，其中 True 值表示该元素应参与注意力，或者一个与查询、键、值具有相同类型的浮点掩码，该掩码被添加到注意力分数中。

过去和现在的键/值状态都是可选的。它们必须一起使用，不允许只使用其中一个。以下模式应用于 Q、K 和 V 输入，并在根据提供的序列长度和头数对 K 和 V 输入进行适当重塑后进行。

The following pattern is applied by this operator:
        Q          K          V
        |          |          |
Q*sqrt(scale) K*sqrt(scale) |
        |          |          |
        |       Transpose     |
        |          |          |
        ---MatMul---          |
            |               |
at_mask---Add              |
            |               |
    softcap (if provided)     |
            |               |
        Softmax            |
            |               |
            -----MatMul------
                    |
                    Y

参数

Q (Tensor) – 查询张量。形状为 (batch_size, q_num_heads, q_sequence_length, head_size) 的 4D 张量，或形状为 (batch_size, q_sequence_length, q_hidden_size) 的 3D 张量。对于 3D 输入张量，q_hidden_size = q_num_heads * head_size
K (Tensor) – 键张量。形状为 (batch_size, kv_num_heads, kv_sequence_length, head_size) 的 4D 张量，或形状为 (batch_size, kv_sequence_length, k_hidden_size) 的 3D 张量。对于 3D 输入张量，k_hidden_size = kv_num_heads * head_size
V (Tensor) – 值张量。形状为 (batch_size, kv_num_heads, kv_sequence_length, v_head_size) 的 4D 张量，或形状为 (batch_size, kv_sequence_length, v_hidden_size) 的 3D 张量。对于 3D 输入张量，v_hidden_size = kv_num_heads * v_head_size
attn_mask (Optional[Tensor]) – 注意力掩码。形状必须可广播到形状为 (batch_size, q_num_heads, q_sequence_length, total_sequence_length) 的 4D 张量，其中 total_sequence_length = past_sequence_length + kv_sequence_length。支持两种类型的掩码。布尔掩码，其中 True 值表示该元素应参与注意力。还支持与查询、键、值具有相同类型的浮点掩码，该掩码被添加到注意力分数中。
past_key (Optional[Tensor]) – 键的 past state 缓存，形状为 (batch_size, kv_num_heads, past_sequence_length, head_size)
past_value (Optional[Tensor]) – 值的 past state 缓存，形状为 (batch_size, kv_num_heads, past_sequence_length, v_head_size)
is_causal (bool) – 如果设置为 True，则当掩码为方形矩阵时，注意力掩码是下三角矩阵。由于对齐，注意力掩码的形式是左上角的因果偏置。
kv_num_heads (int) – 键和值的头数。必须与 Q、K 和 V 的 3D 输入一起使用。
q_num_heads (int) – 查询的头数。必须与 Q、K 和 V 的 3D 输入一起使用。
qk_matmul_output_mode (int) – 如果设置为 0，则 qk_matmul_output 是 qk 乘法的输出。如果设置为 1，则 qk_matmul_output 包括将注意力掩码添加到 qk 乘法输出。如果设置为 2，则 qk_matmul_output 是 softcap 操作后的输出。如果设置为 3，则 qk_matmul_output 是 softmax 操作后的输出。默认值为 0。
scale (Optional[float]) – 应用于 Q*K^T 的缩放因子。默认值为 1/sqrt(head_size)。为防止数值溢出，在乘法之前将 Q、K 乘以 sqrt(scale)。
softcap (float) – 注意力权重的 softcap 值。默认值为 0。
softmax_precision (Optional[int]) – softmax 计算中使用的浮点精度。如果未提供 softmax 精度，则使用与 softmax 输入 (Q 和 K) 相同的精度。

返回

输出张量。形状为 (batch_size, q_num_heads, q_sequence_length, v_head_size) 的 4D 张量，或形状为 (batch_size, q_sequence_length, hidden_size) 的 3D 张量。对于 3D 输入张量，hidden_size = q_num_heads * v_head_size
更新后的键缓存，形状为 (batch_size, kv_num_heads, total_sequence_length, head_size)，其中 total_sequence_length = past_sequence_length + kv_sequence_length。
更新后的值缓存，形状为 (batch_size, kv_num_heads, total_sequence_length, v_head_size)，其中 total_sequence_length = past_sequence_length + kv_sequence_length。
QK 乘法的输出。形状为 (batch_size, q_num_heads, q_sequence_length, total_sequence_length) 的 4D 张量，其中 total_sequence_length = past_sequence_length + kv_sequence_length。

返回类型

一个包含的元组

ONNX 到 ATen 分解表#

您可以使用 torch.onnx.ops.aten_decompositions() 获取一个分解表，将上面定义的 ONNX 运算符分解为 ATen 运算符。

class Model(torch.nn.Module):
    def forward(
        self, input_data, cos_cache_data, sin_cache_data, position_ids_data
    ):
        return torch.onnx.ops.rotary_embedding(
            input_data,
            cos_cache_data,
            sin_cache_data,
            position_ids_data,
        )

model = Model()

ep = torch.export.export(
    model,
    (input_data, cos_cache_data, sin_cache_data, position_ids_data),
)
# The program can be decomposed into aten ops
ep_decomposed = ep.run_decompositions(torch.onnx.ops.aten_decompositions())

torch.onnx.ops.aten_decompositions()[source]#

返回 ONNX 到 ATen 的分解表。

返回类型: dict[torch._ops.OpOverload, Callable]

torch.onnx.ops#

符号运算符#

ONNX 运算符#

ONNX 到 ATen 分解表#

文档

教程

资源