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.ai/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.ai/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 分解表#

文件

教程

資源