基於 TorchScript 的 ONNX 匯出器¶
注意
若要使用 TorchDynamo（而不是 TorchScript）匯出 ONNX 模型，請參閱 torch.onnx.dynamo_export()。
示例：將 AlexNet 從 PyTorch 轉換為 ONNX¶
這是一個簡單的指令碼，用於將預訓練的 AlexNet 匯出到名為 alexnet.onnx 的 ONNX 檔案。呼叫 torch.onnx.export 一次執行模型以跟蹤其執行情況，然後將跟蹤的模型匯出到指定的檔案
import torch
import torchvision
dummy_input = torch.randn(10, 3, 224, 224, device="cuda")
model = torchvision.models.alexnet(pretrained=True).cuda()
# Providing input and output names sets the display names for values
# within the model's graph. Setting these does not change the semantics
# of the graph; it is only for readability.
#
# The inputs to the network consist of the flat list of inputs (i.e.
# the values you would pass to the forward() method) followed by the
# flat list of parameters. You can partially specify names, i.e. provide
# a list here shorter than the number of inputs to the model, and we will
# only set that subset of names, starting from the beginning.
input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]
torch.onnx.export(model, dummy_input, "alexnet.onnx", verbose=True, input_names=input_names, output_names=output_names)
生成的 alexnet.onnx 檔案包含一個二進位制 協議緩衝區，其中包含您匯出的模型（在本例中為 AlexNet）的網路結構和引數。引數 verbose=True 會導致匯出器打印出模型的可讀表示
# These are the inputs and parameters to the network, which have taken on
# the names we specified earlier.
graph(%actual_input_1 : Float(10, 3, 224, 224)
%learned_0 : Float(64, 3, 11, 11)
%learned_1 : Float(64)
%learned_2 : Float(192, 64, 5, 5)
%learned_3 : Float(192)
# ---- omitted for brevity ----
%learned_14 : Float(1000, 4096)
%learned_15 : Float(1000)) {
# Every statement consists of some output tensors (and their types),
# the operator to be run (with its attributes, e.g., kernels, strides,
# etc.), its input tensors (%actual_input_1, %learned_0, %learned_1)
%17 : Float(10, 64, 55, 55) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[11, 11], pads=[2, 2, 2, 2], strides=[4, 4]](%actual_input_1, %learned_0, %learned_1), scope: AlexNet/Sequential[features]/Conv2d[0]
%18 : Float(10, 64, 55, 55) = onnx::Relu(%17), scope: AlexNet/Sequential[features]/ReLU[1]
%19 : Float(10, 64, 27, 27) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%18), scope: AlexNet/Sequential[features]/MaxPool2d[2]
# ---- omitted for brevity ----
%29 : Float(10, 256, 6, 6) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%28), scope: AlexNet/Sequential[features]/MaxPool2d[12]
# Dynamic means that the shape is not known. This may be because of a
# limitation of our implementation (which we would like to fix in a
# future release) or shapes which are truly dynamic.
%30 : Dynamic = onnx::Shape(%29), scope: AlexNet
%31 : Dynamic = onnx::Slice[axes=[0], ends=[1], starts=[0]](%30), scope: AlexNet
%32 : Long() = onnx::Squeeze[axes=[0]](%31), scope: AlexNet
%33 : Long() = onnx::Constant[value={9216}](), scope: AlexNet
# ---- omitted for brevity ----
%output1 : Float(10, 1000) = onnx::Gemm[alpha=1, beta=1, broadcast=1, transB=1](%45, %learned_14, %learned_15), scope: AlexNet/Sequential[classifier]/Linear[6]
return (%output1);
}
您還可以使用 ONNX 庫來驗證輸出，您可以使用 pip 安裝該庫
pip install onnx
然後，您可以執行
import onnx
# Load the ONNX model
model = onnx.load("alexnet.onnx")
# Check that the model is well formed
onnx.checker.check_model(model)
# Print a human readable representation of the graph
print(onnx.helper.printable_graph(model.graph))
您還可以使用眾多支援 ONNX 的 執行時 中的一種來執行匯出的模型。例如，在安裝 ONNX Runtime 後，您可以載入並執行模型
import onnxruntime as ort
import numpy as np
ort_session = ort.InferenceSession("alexnet.onnx")
outputs = ort_session.run(
None,
{"actual_input_1": np.random.randn(10, 3, 224, 224).astype(np.float32)},
)
print(outputs[0])
這是一個更復雜的 教程，介紹如何匯出模型並使用 ONNX Runtime 執行它。
跟蹤與指令碼編寫¶
從內部來看，torch.onnx.export() 要求 torch.jit.ScriptModule 而不是 torch.nn.Module。如果傳入的模型不是 ScriptModule，export() 會使用追蹤功能將其轉換為 ScriptModule
追蹤：如果
torch.onnx.export()與不是ScriptModule的模組一起呼叫，那麼它會首先執行相當於torch.jit.trace()的操作，即使用給定的args執行模型一次，並記錄執行期間發生的所有操作。這意味著如果模型是動態的，例如根據輸入資料改變行為，那麼匯出的模型不會捕獲此動態行為。我們建議檢查匯出的模型，並確保運算子看起來合理。追蹤會展開迴圈和 if 語句，並匯出與追蹤執行完全相同的靜態圖。如果想要動態控制流程地匯出模型，則需要使用指令碼編制。指令碼編制：透過指令碼編制編譯模型可以保留動態控制流程，並且對於不同大小的輸入有效。要使用指令碼編制，請執行以下操作
使用
torch.jit.script()生成ScriptModule。使用
ScriptModule作為模型呼叫torch.onnx.export()。args仍然是必需的，但它們僅在內部用於生成示例輸出，以便可以捕獲輸出的型別和形狀。不會執行追蹤。
有關更多詳細資訊，包括如何整合追蹤和指令碼編制以滿足不同模型的特定需求，請參閱 TorchScript 簡介 和 TorchScript。
避免陷入誤區¶
避免使用 NumPy 和內建 Python 型別¶
PyTorch 模型可以使用 NumPy 或 Python 型別和函式編寫，但在 跟蹤 期間，任何 NumPy 或 Python 型別（而非 torch.Tensor）的變數都會轉換為常量，如果這些值應隨輸入而更改，這將產生錯誤的結果。
例如，不要在 numpy.ndarrays 上使用 numpy 函式，而要
# Bad! Will be replaced with constants during tracing.
x, y = np.random.rand(1, 2), np.random.rand(1, 2)
np.concatenate((x, y), axis=1)
在 torch.Tensors 上使用 torch 運算子
# Good! Tensor operations will be captured during tracing.
x, y = torch.randn(1, 2), torch.randn(1, 2)
torch.cat((x, y), dim=1)
不要使用 torch.Tensor.item()（它將張量轉換為 Python 內建數字），而要
# Bad! y.item() will be replaced with a constant during tracing.
def forward(self, x, y):
return x.reshape(y.item(), -1)
使用 torch 對單元素張量的隱式轉換的支援
# Good! y will be preserved as a variable during tracing.
def forward(self, x, y):
return x.reshape(y, -1)
避免使用 Tensor.data¶
使用 Tensor.data 欄位可能會產生錯誤的跟蹤資訊，進而產生錯誤的 ONNX 圖。請改用 torch.Tensor.detach()。（正在進行移除 Tensor.data 的工作，請在此處檢視）。
在追蹤模式下使用 tensor.shape 時避免就地操作¶
在追蹤模式下，從 tensor.shape 獲取的形狀被追蹤為張量，並共享相同的記憶體。這可能會導致最終輸出值不匹配。作為一種解決方法，在此類場景中避免使用就地操作。例如，在模型中
class Model(torch.nn.Module):
def forward(self, states):
batch_size, seq_length = states.shape[:2]
real_seq_length = seq_length
real_seq_length += 2
return real_seq_length + seq_length
real_seq_length 和 seq_length 在追蹤模式下共享相同的記憶體。這可以透過重寫就地操作來避免
real_seq_length = real_seq_length + 2
侷限性¶
型別¶
僅
torch.Tensors、可以輕鬆轉換為 torch.Tensors 的數字型別（如 float、int）以及這些型別的元組和列表受支援作為模型輸入或輸出。字典和 str 輸入和輸出在 跟蹤 模式下被接受，但任何依賴於詞典或 str 輸入值的計算都將用在一條追蹤執行期間看到的常數值替換。
任何輸出是字典的都會被靜默地替換為其值的扁平序列（鍵將被移除）。例如：
{"foo": 1, "bar": 2}將變成(1, 2)。任何輸出是 str 的都會被靜默地移除。
由於 ONNX 對巢狀序列的支援有限，因此在 指令碼 模式中不支援涉及元組和列表的某些操作。特別是並不支援將元組追加到列表中。在跟蹤模式中，將在跟蹤期間自動使巢狀序列變平。
運算元實現中的差異¶
由於運算元實現的不同，在不同的執行時上執行匯出模型可能會產生彼此或與 PyTorch 不同的結果。通常情況下，這些差異在數值上非常小，所以只有當您的應用程式對此類微小差異很敏感時，這才是您需要擔心的問題。
不受支援的張量索引模式¶
以下是無法匯出的張量索引模式的列表。如果您正在匯出不包含任何以下不受支援模式的模型時遇到問題，請仔細檢查您是否使用最新的 opset_version 進行匯出。
讀取/獲取¶
在對張量進行索引以進行讀取時，不支援以下模式
# Tensor indices that includes negative values.
data[torch.tensor([[1, 2], [2, -3]]), torch.tensor([-2, 3])]
# Workarounds: use positive index values.
寫入/設定¶
在對張量進行索引以進行寫入時，不支援以下模式
# Multiple tensor indices if any has rank >= 2
data[torch.tensor([[1, 2], [2, 3]]), torch.tensor([2, 3])] = new_data
# Workarounds: use single tensor index with rank >= 2,
# or multiple consecutive tensor indices with rank == 1.
# Multiple tensor indices that are not consecutive
data[torch.tensor([2, 3]), :, torch.tensor([1, 2])] = new_data
# Workarounds: transpose `data` such that tensor indices are consecutive.
# Tensor indices that includes negative values.
data[torch.tensor([1, -2]), torch.tensor([-2, 3])] = new_data
# Workarounds: use positive index values.
# Implicit broadcasting required for new_data.
data[torch.tensor([[0, 2], [1, 1]]), 1:3] = new_data
# Workarounds: expand new_data explicitly.
# Example:
# data shape: [3, 4, 5]
# new_data shape: [5]
# expected new_data shape after broadcasting: [2, 2, 2, 5]
新增對運算元的支援¶
在匯出包含不受支援運算元的模型時，您將會看到這樣的錯誤訊息
RuntimeError: ONNX export failed: Couldn't export operator foo
當出現這種情況時，您可以執行以下幾項操作
更改模型以避免使用該運算元。
建立一個符號函式來轉換運算元並將其註冊為自定義符號函式。
為 PyTorch 作出貢獻，將同一個符號函式新增到
torch.onnx本身。
如果您決定實現一個符號函式（我們希望您能將它貢獻回 PyTorch！），請使用以下方法開始
ONNX 匯出器內部¶
“符號函式”是一個函式，它將 PyTorch 運算元分解為一系列 ONNX 運算元的組合。
在匯出期間，匯出器將按拓撲順序訪問 TorchScript 圖中的每個節點（包含一個 PyTorch 運算元）。在訪問一個節點時，匯出器將查詢針對該運算元的已註冊的符號函式。符號函式在 Python 中實現。名為 foo 的運算元的符號函式看起來類似於
def foo(
g,
input_0: torch._C.Value,
input_1: torch._C.Value) -> Union[None, torch._C.Value, List[torch._C.Value]]:
"""
Adds the ONNX operations representing this PyTorch function by updating the
graph g with `g.op()` calls.
Args:
g (Graph): graph to write the ONNX representation into.
input_0 (Value): value representing the variables which contain
the first input for this operator.
input_1 (Value): value representing the variables which contain
the second input for this operator.
Returns:
A Value or List of Values specifying the ONNX nodes that compute something
equivalent to the original PyTorch operator with the given inputs.
None if it cannot be converted to ONNX.
"""
...
torch._C 型別是 C++ 中定義的型別的 Python 包裝器，這些型別在 ir.h 中定義。
新增符號函式的過程取決於運算元型別。
ATen 運算元¶
ATen 是 PyTorch 的內部張量庫。如果該運算元是 ATen 運算元（在 TorchScript 圖中以字首 aten:: 出現），確保還沒有支援它。
支援的運算元列表¶
訪問自動生成的 支援的 TorchScript 運算元列表，詳細瞭解每個 opset_version 中支援哪些運算元。
為一個 aten 或量化運算元新增支援¶
如果運算元沒有出現在上述列表中
在
torch/onnx/symbolic_opset<version>.py中定義符號函式，例如 torch/onnx/symbolic_opset9.py”。確保此函式具有與 ATen 函式相同的名字，該函式可能在torch/_C/_VariableFunctions.pyi或torch/nn/functional.pyi中宣告（這些檔案在生成時生成，因此在你構建 PyTorch 之前，這些檔案不會出現在你的簽出中）。預設情況下，第一個引數是 ONNX 圖。其他引數名必須嚴格匹配
.pyi檔案中的名稱，因為分派使用關鍵字引數完成。在符號函式中，如果運算元在 ONNX 標準運算元集 中，我們只需要建立一個節點來表示圖中的 ONNX 運算元。如果不這樣做，我們可以組成幾個具有與 ATen 運算元等效語義的標準運算元。
下面是處理 ELU 運算元的缺失符號函式的一個示例。
如果我們執行以下程式碼
print(
torch.jit.trace(
torch.nn.ELU(), # module
torch.ones(1) # example input
).graph
)
我們會看到類似於
graph(%self : __torch__.torch.nn.modules.activation.___torch_mangle_0.ELU,
%input : Float(1, strides=[1], requires_grad=0, device=cpu)):
%4 : float = prim::Constant[value=1.]()
%5 : int = prim::Constant[value=1]()
%6 : int = prim::Constant[value=1]()
%7 : Float(1, strides=[1], requires_grad=0, device=cpu) = aten::elu(%input, %4, %5, %6)
return (%7)
由於我們在圖中看到了 aten::elu，所以我們知道這是一個 ATen 運算子。
我們檢視 ONNX 運算子列表，並確認 Elu 在 ONNX 中是標準化的。
我們在 torch/nn/functional.pyi 中找到了 elu 的簽名
def elu(input: Tensor, alpha: float = ..., inplace: bool = ...) -> Tensor: ...
我們將下列行新增到 symbolic_opset9.py
def elu(g, input: torch.Value, alpha: torch.Value, inplace: bool = False):
return g.op("Elu", input, alpha_f=alpha)
現在 PyTorch 能夠匯出包含 aten::elu 運算子的模型！
請參閱 torch/onnx/symbolic_opset*.py 檔案以瞭解更多示例。
torch.autograd.Functions¶
如果運算子是 torch.autograd.Function 的子類，則有三種方法可以匯出它。
靜態符號方法¶
可以向你的函式類新增一個名為 symbolic 的靜態方法。它應該返回在 ONNX 中表示函式行為的 ONNX 運算子。例如
class MyRelu(torch.autograd.Function):
@staticmethod
def forward(ctx, input: torch.Tensor) -> torch.Tensor:
ctx.save_for_backward(input)
return input.clamp(min=0)
@staticmethod
def symbolic(g: torch.Graph, input: torch.Value) -> torch.Value:
return g.op("Clip", input, g.op("Constant", value_t=torch.tensor(0, dtype=torch.float)))
內聯自動微分函式¶
在針對其後續 torch.autograd.Function 未提供靜態符號方法或者在未提供將 prim::PythonOp 註冊為自定義符號函式的函式的情況下，torch.onnx.export() 會嘗試內聯對應於該 torch.autograd.Function 的圖，以便將此函式分解為函式中使用的單獨運算子。只要支援這些單獨的運算子，匯出就應成功。例如
class MyLogExp(torch.autograd.Function):
@staticmethod
def forward(ctx, input: torch.Tensor) -> torch.Tensor:
ctx.save_for_backward(input)
h = input.exp()
return h.log().log()
此模型沒有出現靜態符號方法，但會按如下方式匯出
graph(%input : Float(1, strides=[1], requires_grad=0, device=cpu)):
%1 : float = onnx::Exp[](%input)
%2 : float = onnx::Log[](%1)
%3 : float = onnx::Log[](%2)
return (%3)
如果您需要避免內聯torch.autograd.Function，應將模型匯出為 operator_export_type 設定為 ONNX_FALLTHROUGH 或 ONNX_ATEN_FALLBACK。
自定義運算元¶
您可以匯出包含多種標準 ONNX op 組合的模型或由自定義 C++ 後端驅動模型。
ONNX-script 函式¶
如果運算元不是標準 ONNX op，但可以由多個現有的 ONNX op 組成，您可以利用ONNX-script建立一個外部 ONNX 函式來支援該運算元。您可以按照以下示例進行匯出
import onnxscript
# There are three opset version needed to be aligned
# This is (1) the opset version in ONNX function
from onnxscript.onnx_opset import opset15 as op
opset_version = 15
x = torch.randn(1, 2, 3, 4, requires_grad=True)
model = torch.nn.SELU()
custom_opset = onnxscript.values.Opset(domain="onnx-script", version=1)
@onnxscript.script(custom_opset)
def Selu(X):
alpha = 1.67326 # auto wrapped as Constants
gamma = 1.0507
alphaX = op.CastLike(alpha, X)
gammaX = op.CastLike(gamma, X)
neg = gammaX * (alphaX * op.Exp(X) - alphaX)
pos = gammaX * X
zero = op.CastLike(0, X)
return op.Where(X <= zero, neg, pos)
# setType API provides shape/type to ONNX shape/type inference
def custom_selu(g: jit_utils.GraphContext, X):
return g.onnxscript_op(Selu, X).setType(X.type())
# Register custom symbolic function
# There are three opset version needed to be aligned
# This is (2) the opset version in registry
torch.onnx.register_custom_op_symbolic(
symbolic_name="aten::selu",
symbolic_fn=custom_selu,
opset_version=opset_version,
)
# There are three opset version needed to be aligned
# This is (2) the opset version in exporter
torch.onnx.export(
model,
x,
"model.onnx",
opset_version=opset_version,
# only needed if you want to specify an opset version > 1.
custom_opsets={"onnx-script": 2}
)
上述示例將其匯出為“onnx-script”opset 中的自定義運算元。匯出自定義運算元時，可以使用 custom_opsets 詞典在匯出時指定自定義域版本。如果未指定，自定義 opset 版本預設為 1。
注意：請小心對齊上述示例中提到的 opset 版本，並確保在匯出步驟中使用它們。如何編寫 onnx-script 函式的示例用法是 onnx-script 中積極開發的一個 beta 版本。請遵循最新的ONNX-script
C++ 運算元¶
如果模型使用以 C++ 實現的自定義運算元（如使用自定義 C++ 運算元擴充套件 TorchScript中所述），您可以按照以下示例進行匯出
from torch.onnx import symbolic_helper
# Define custom symbolic function
@symbolic_helper.parse_args("v", "v", "f", "i")
def symbolic_foo_forward(g, input1, input2, attr1, attr2):
return g.op("custom_domain::Foo", input1, input2, attr1_f=attr1, attr2_i=attr2)
# Register custom symbolic function
torch.onnx.register_custom_op_symbolic("custom_ops::foo_forward", symbolic_foo_forward, 9)
class FooModel(torch.nn.Module):
def __init__(self, attr1, attr2):
super().__init__()
self.attr1 = attr1
self.attr2 = attr2
def forward(self, input1, input2):
# Calling custom op
return torch.ops.custom_ops.foo_forward(input1, input2, self.attr1, self.attr2)
model = FooModel(attr1, attr2)
torch.onnx.export(
model,
(example_input1, example_input1),
"model.onnx",
# only needed if you want to specify an opset version > 1.
custom_opsets={"custom_domain": 2}
)
上述示例將其匯出為“custom_domain”opset 中的自定義運算元。匯出自定義運算元時，可以使用 custom_opsets 詞典在匯出時指定自定義域版本。如果未指定，自定義 opset 版本預設為 1。
使用該模型的執行時需要支援自定義 op。請參閱Caffe2 自定義 op、ONNX Runtime 自定義 op或所選執行時的文件。
一次發現所有不可轉換的 Aten op¶
如果匯出因無法轉換的 ATen op 而失敗，實際上可能存在多於一個這樣的 op，但錯誤訊息僅提及第一個 op。你可以一次性發現所有無法轉換的 op
# prepare model, args, opset_version
...
torch_script_graph, unconvertible_ops = torch.onnx.utils.unconvertible_ops(
model, args, opset_version=opset_version
)
print(set(unconvertible_ops))
該集合是個近似，因為某些 op 可能會在轉換過程中被移除且無需轉換。某些其他 op 可能具備部分支援，但使用特定輸入時會轉換失敗，但這應能讓你大致瞭解不支援哪些 op。歡迎在 GitHub 上提出 op 支援請求。
常見問題¶
問：我已匯出 LSTM 模型，但其輸入大小似乎已固定？
追蹤器會記錄示例輸入的形狀。如果模型應接受動態形狀的輸入，請在呼叫
torch.onnx.export()時設定dynamic_axes。
問：如何匯出包含迴圈的模型？
請參閱 追蹤與指令碼。
問：如何匯出帶有原始型別輸入（如 int、float）的模型？
在 PyTorch 1.9 中添加了對原始數字型別輸入的支援。但是，匯出器不支援帶有 str 輸入的模型。
問：ONNX 是否支援隱式標量資料型別強制轉換？
ONNX 標準不支援，但匯出器會嘗試處理該部分。標量以常量張量形式匯出。匯出器會找出標量正確的型別。在少數無法完成此操作的情況下，你需要手動用 dtype=torch.float32 等指定資料型別。如果你看到任何錯誤，請 [建立 GitHub issue](https://github.com/pytorch/pytorch/issues).
問：張量列表可以匯出到 ONNX 嗎？
可以，對於
opset_version>= 11，因為 ONNX 在 opset 11 中引入了 Sequence 型別。
Python API¶
函式¶
- torch.onnx.export(model, args, f=None, export_params=True, verbose=False, training=<TrainingMode.EVAL: 0>, input_names=None, output_names=None, operator_export_type=<OperatorExportTypes.ONNX: 0>, opset_version=None, do_constant_folding=True, dynamic_axes=None, keep_initializers_as_inputs=None, custom_opsets=None, export_modules_as_functions=False, autograd_inlining=True, dynamo=False)[source]¶
以 ONNX 格式匯出模型。
如果
model不是torch.jit.ScriptModule或torch.jit.ScriptFunction，這會執行model一次，以便將其轉換為要匯出的 TorchScript 圖（相當於torch.jit.trace()）。因此，它對動態控制流的支援與torch.jit.trace()的支援相同。- 引數
model (
torch.nn.Module,torch.jit.ScriptModule或torch.jit.ScriptFunction) –要匯出的模型。args (元組 或 torch.Tensor) –
args 可以構成如下形式：
僅引數元組
args = (x, y, z)
元組應包含模型輸入，確保
model(*args)是對模型的有效呼叫。任何非張量引數都將硬編碼到匯出的模型中；任何張量引數都將按其在元組中出現的順序成為匯出模型的輸入。張量
args = torch.Tensor([1])
這等效於該張量的 1 元組。
以命名引數字典結尾的引數元組
args = ( x, { "y": input_y, "z": input_z } )
元組的所有元素都作為非關鍵字引數傳遞，而命名引數將從最後一個元素中設定。如果命名引數不存在於字典中，則會向其分配預設值，如果沒有提供預設值，則該值為空。
注意
如果字典是 args 元組的最後一個元素，則會在其中解釋為包含命名引數。為了將字典作為最後一個非關鍵字引數傳遞，請提供一個空字典作為 args 元組的最後一個元素。例如，請執行以下操作，而不是
torch.onnx.export( model, ( x, # WRONG: will be interpreted as named arguments {y: z} ), "test.onnx.pb" )
編寫以下內容
torch.onnx.export( model, ( x, {y: z}, {} ), "test.onnx.pb" )
f (可選項[並集[str, BytesIO]]) - 檔案的類似物件（此類物件
f.fileno()會返回檔案描述符）或包含檔名字串。二進位制協議緩衝區將被寫入此檔案。export_params (布林值, 預設值為 True) - 如果為 True，則將匯出所有引數。如果您想匯出未訓練的模型，請將其設定為 False。在這種情況下，匯出的模型將首先獲取其所有引數作為引數，排序則根據
model.state_dict().values()指定的先後順序verbose (布林值, 預設值為 False) - 如果為 True，則會列印要匯出到 stdout 的模型說明。此外，最終的 ONNX 圖表將包含匯出模型中
doc_string`域，其中會提及model的原始碼位置。如果為 True，則將啟動 ONNX 匯出程式記錄。training (列舉, 預設值為 TrainingMode.EVAL) -
TrainingMode.EVAL：以推理模式匯出模型。TrainingMode.PRESERVE：如果 model.training 為False，則以推理模式匯出模型，如果 model.training 為 True，則以訓練模式匯出模型。
TrainingMode.TRAINING：以訓練模式匯出模型。停用最佳化可能會干擾訓練。
operator_export_type (列舉, 預設 OperatorExportTypes.ONNX) -
OperatorExportTypes.ONNX：將所有運算子匯出為常規 ONNX 運算子（在預設的 opset 域中）。
OperatorExportTypes.ONNX_FALLTHROUGH：嘗試轉換所有運算子到預設 opset 域中的標準 ONNX 運算子。如果無法這樣做（例如，因為尚未新增將特定 PyTorch 運算子轉換為 ONNX 的支援），則退回到將運算子匯出到自定義 opset 域而不進行轉換。適用於 自定義運算子 和 ATen 運算子。為了使匯出的模型可用，執行時必須支援這些非標準運算子。
OperatorExportTypes.ONNX_ATEN：所有 ATen 運算子（在 TorchScript 名稱空間 “aten” 中）被匯出為 ATen 運算子（在 opset 域 “org.pytorch.aten” 中）。ATen 是 PyTorch 的內建張量庫，所以這會指示執行時使用 PyTorch 的這些運算子的實現。
警告
透過這種方式匯出的模型可能只能由 Caffe2 執行。
如果運算元實現中的數值差異導致 PyTorch 和 Caffe2 之間存在較大的行為差異（在未訓練的模型中更為常見），這可能很有用。
OperatorExportTypes.ONNX_ATEN_FALLBACK：嘗試匯出每個 ATen op(在 TorchScript 名稱空間 “aten” 中) 作為常規 ONNX op。如果我們無法這樣做（例如，因為尚未新增將特定 PyTorch 運算子轉換為 ONNX 的支援），則退回到匯出 ATen op。請參閱 OperatorExportTypes.ONNX_ATEN 的文件以瞭解上下文。例如
graph(%0 : Float): %3 : int = prim::Constant[value=0]() # conversion unsupported %4 : Float = aten::triu(%0, %3) # conversion supported %5 : Float = aten::mul(%4, %0) return (%5)
假設
aten::triu在 ONNX 中不受支援，這將匯出為graph(%0 : Float): %1 : Long() = onnx::Constant[value={0}]() # not converted %2 : Float = aten::ATen[operator="triu"](%0, %1) # converted %3 : Float = onnx::Mul(%2, %0) return (%3)
警告
透過這種方式匯出的模型可能只能由 Caffe2 執行。
opset_version (整數, 預設 17) - 預設的 (ai.onnx) opset 的版本的目標。必須大於等於 7 且小於等於 17。
do_constant_folding (bool, 預設值 True) – 應用常量摺疊最佳化。常量摺疊將用預先計算好的常量節點替換掉所有常量輸入的某些 op。
dynamic_axes (dict[string, dict[int, string]] 或 dict[string, list(int)], 預設值為空字典) –
預設情況下，匯出的模型將具有所有輸入和輸出張量的形狀，確切匹配在
args中給定的形狀。如需將張量的軸指定為動態（即在執行時才知道），請將dynamic_axes設定為符合架構的字典- 鍵 (str)：輸入或輸出名稱。每個名稱也必須提供在
input_names或 output_names 中.
- 鍵 (str)：輸入或輸出名稱。每個名稱也必須提供在
- 值 (字典或列表)：如果為字典，則鍵為軸索引，值為軸名稱。如果為
列表，每個元素都是軸索引。
例如
class SumModule(torch.nn.Module): def forward(self, x): return torch.sum(x, dim=1) torch.onnx.export( SumModule(), (torch.ones(2, 2),), "onnx.pb", input_names=["x"], output_names=["sum"] )
產生
input { name: "x" ... shape { dim { dim_value: 2 # axis 0 } dim { dim_value: 2 # axis 1 ... output { name: "sum" ... shape { dim { dim_value: 2 # axis 0 ...
當
torch.onnx.export( SumModule(), (torch.ones(2, 2),), "onnx.pb", input_names=["x"], output_names=["sum"], dynamic_axes={ # dict value: manually named axes "x": {0: "my_custom_axis_name"}, # list value: automatic names "sum": [0], } )
產生
input { name: "x" ... shape { dim { dim_param: "my_custom_axis_name" # axis 0 } dim { dim_value: 2 # axis 1 ... output { name: "sum" ... shape { dim { dim_param: "sum_dynamic_axes_1" # axis 0 ...
keep_initializers_as_inputs (bool, 預設值 None) –
如果為 True，則匯出的圖中的所有初始化器（通常對應於引數）也將作為輸入新增到圖中。如果為 False，則不會將初始化器作為輸入新增到圖中，僅將非引數輸入作為輸入新增。這可能允許後端/執行時進行更好的最佳化（例如常量摺疊）。
如果為 True，則不會執行 deduplicate_initializers 傳遞。這意味著具有重複值的初始化器不會進行重複資料刪除，並會被視為圖中的不同輸入。這允許在匯出後在執行時提供不同的輸入初始化器。
如果
opset_version < 9，則初始化器必須是圖輸入的一部分，並且將忽略此引數，並且行為等同於將此引數設定為 True。如果為 None，則自動選擇行為，如下所示
- 如果
operator_export_type=OperatorExportTypes.ONNX，則行為等同於 將此引數設定為 False。
- 如果
其他情況下，行為等同於將此引數設定為 True。
custom_opsets (dict[str, int], 預設 empty dict) –
具有架構的字典
鍵 (str): opset 域名稱
值 (int): opset 版本
如果
model引用了自定義 opset，但此字典中未提及，則 opset 版本設定為 1。只能透過此引數指示自定義 opset 域名稱和版本。export_modules_as_functions (bool 或 set of type of nn.Module, 預設值 False) –
啟用將所有
nn.Module轉發呼叫作為 ONNX 中的區域性函式匯出的標誌。或指示要作為 ONNX 中區域性函式匯出的模組型別。此功能需要opset_version>= 15，否則匯出會失敗。這是因為opset_version< 15 意味著 IR 版本 < 8，這意味著不支援區域性函式。模組變數將作為函式屬性匯出。共有兩類函式屬性。1. 帶註釋的屬性：透過 PEP 526 樣式 進行型別註釋的類變數將作為屬性匯出。ONNX 區域性函式的子圖內未使用帶註釋的屬性，因為它們不是由 PyTorch JIT 跟蹤建立的，但使用者可能會使用它們來確定是否使用特定融合核替換函式。
2. 推斷出的屬性：模組中運算子使用的變數。屬性名稱將帶有字首“inferred::”。這是為了區別於從 Python 模組註釋中檢索到的預定義屬性。推斷出的屬性用於 ONNX 區域性函式的子圖內。
False（預設值）：將nn.Module轉發呼叫匯出為細粒度的節點。True：將所有nn.Module轉發呼叫匯出為區域性函式節點。- 型別
nn.Module的集合：匯出nn.Moduleforward 呼叫為區域性函式節點， 僅當該
nn.Module的型別存在於該集合中時。
- 型別
autograd_inlining (布林值，預設為 True) - 用於控制是否內聯自動梯度函式的標誌。有關更多詳細資訊，請參閱https://github.com/pytorch/pytorch/pull/74765。
dynamo (布林值，預設為 False) - 是否使用 Dynamo（而不是 TorchScript）匯出模型。
- 引發
torch.onnx.errors.CheckerError - 如果 ONNX 檢查器檢測到無效的 ONNX 圖形。
torch.onnx.errors.UnsupportedOperatorError - 如果 ONNX 圖形因使用匯出器不支援的運算子而無法匯出。
torch.onnx.errors.OnnxExporterError - 匯出過程中可能發生的其他錯誤。所有錯誤均為
errors.OnnxExporterError的子類。
- 返回型別
- torch.onnx.export_to_pretty_string(model, args, export_params=True, verbose=False, training=<TrainingMode.EVAL: 0>, input_names=None, output_names=None, operator_export_type=<OperatorExportTypes.ONNX: 0>, export_type=None, google_printer=False, opset_version=None, keep_initializers_as_inputs=None, custom_opsets=None, add_node_names=True, do_constant_folding=True, dynamic_axes=None)[原始碼]¶
類似於
export()，但返回 ONNX 模型的文字表示。僅在下面列出引數時存在差異。所有其他引數都與export()相同。
- torch.onnx.register_custom_op_symbolic(symbolic_name, symbolic_fn, opset_version)[source]¶
為自定義運算子註冊符號函式。
當用戶為自定義/contrib 運算子註冊符號時強烈建議透過 setType API 為該運算子新增形狀推斷，否則在一些極端情況下匯出的圖可能具有不正確的形狀推斷。setType 的示例是 test_operators.py 中的 test_aten_embedding_2。
有關示例用法，請參見模組文件中的“自定義運算子”。
- torch.onnx.unregister_custom_op_symbolic(symbolic_name, opset_version)[原始碼]¶
登出
symbolic_name。有關示例用法，請參見模組文件中的“自定義運算子”。
- torch.onnx.select_model_mode_for_export(model, mode)[原始碼]¶
一個上下文管理器，用於暫時設定
model的訓練模式為mode，並當離開 with-block 時將其重置。
- torch.onnx.verification.find_mismatch(model, input_args, do_constant_folding=True, training=<TrainingMode.EVAL: 0>, opset_version=None, keep_initializers_as_inputs=True, verbose=False, options=None)[source]¶
找出原始模型和匯出模型之間的所有不匹配項。
實驗性。此 API 會發生更改。
此工具有助於除錯原始 PyTorch 模型和匯出的 ONNX 模型之間的不匹配。它對模型圖進行二分查詢，以找出出現不匹配的最小子圖。
- 引數
model (Union[Module, ScriptModule]) – 要匯出的模型。
do_constant_folding (bool) – 與
torch.onnx.export()中的 do_constant_folding 相同。training (TrainingMode) – 與
torch.onnx.export()中的 training 相同。opset_version (Optional[int]) – 與
torch.onnx.export()中的 opset_version 相同。keep_initializers_as_inputs (bool) – 與
torch.onnx.export()中的 keep_initializers_as_inputs 相同。verbose (bool) – 與
torch.onnx.export()中的 verbose 相同。options (Optional[VerificationOptions]) – 驗證不匹配的選項。
- 返回
一個包含不匹配資訊 GraphInfo 物件。
- 返回型別
例項
>>> import torch >>> import torch.onnx.verification >>> torch.manual_seed(0) >>> opset_version = 15 >>> # Define a custom symbolic function for aten::relu. >>> # The custom symbolic function is incorrect, which will result in mismatches. >>> def incorrect_relu_symbolic_function(g, self): ... return self >>> torch.onnx.register_custom_op_symbolic( ... "aten::relu", ... incorrect_relu_symbolic_function, ... opset_version=opset_version, ... ) >>> class Model(torch.nn.Module): ... def __init__(self): ... super().__init__() ... self.layers = torch.nn.Sequential( ... torch.nn.Linear(3, 4), ... torch.nn.ReLU(), ... torch.nn.Linear(4, 5), ... torch.nn.ReLU(), ... torch.nn.Linear(5, 6), ... ) ... def forward(self, x): ... return self.layers(x) >>> graph_info = torch.onnx.verification.find_mismatch( ... Model(), ... (torch.randn(2, 3),), ... opset_version=opset_version, ... ) ===================== Mismatch info for graph partition : ====================== ================================ Mismatch error ================================ Tensor-likes are not close! Mismatched elements: 12 / 12 (100.0%) Greatest absolute difference: 0.2328854203224182 at index (1, 2) (up to 1e-07 allowed) Greatest relative difference: 0.699536174352349 at index (1, 3) (up to 0.001 allowed) ==================================== Tree: ===================================== 5 X __2 X __1 \u2713 id: | id: 0 | id: 00 | | | |__1 X (aten::relu) | id: 01 | |__3 X __1 \u2713 id: 1 | id: 10 | |__2 X __1 X (aten::relu) id: 11 | id: 110 | |__1 \u2713 id: 111 =========================== Mismatch leaf subgraphs: =========================== ['01', '110'] ============================= Mismatch node kinds: ============================= {'aten::relu': 2}
