注意
點擊這裡下載完整的範例程式碼
ONNX 簡介 || 將 PyTorch 模型匯出至 ONNX || 擴展 ONNX 註冊表
將 PyTorch 模型匯出至 ONNX¶
建立於：2023 年 10 月 04 日 | 最後更新：2025 年 1 月 27 日 | 最後驗證：2024 年 11 月 05 日
作者: Ti-Tai Wang 和 Xavier Dupré
注意
截至 PyTorch 2.1，有兩個版本的 ONNX 匯出器。
torch.onnx.dynamo_export是最新的（仍在 beta 版）匯出器，基於 PyTorch 2.0 隨附的 TorchDynamo 技術torch.onnx.export基於 TorchScript 後端，自 PyTorch 1.2.0 起可用
在 60 分鐘速成班中，我們有機會在高階層面了解 PyTorch，並訓練一個小型神經網路來對圖像進行分類。在本教學中，我們將擴展此內容，描述如何使用 TorchDynamo 和 torch.onnx.dynamo_export ONNX 匯出器將 PyTorch 中定義的模型轉換為 ONNX 格式。
雖然 PyTorch 非常適合迭代開發模型，但可以使用不同的格式將模型部署到生產環境，包括 ONNX (開放神經網路交換)！
ONNX 是一種靈活的開放標準格式，用於表示機器學習模型，機器學習的標準化表示允許它們在各種硬體平台和執行時間環境中執行，從大型雲端超級電腦到資源受限的邊緣裝置，例如您的網頁瀏覽器和手機。
在本教學中，我們將學習如何
安裝所需的依賴項。
建立一個簡單的圖像分類器模型。
將模型匯出為 ONNX 格式。
將 ONNX 模型儲存在檔案中。
使用 Netron 可視化 ONNX 模型圖。
使用 ONNX Runtime 執行 ONNX 模型
將 PyTorch 結果與 ONNX Runtime 的結果進行比較。
1. 安裝所需的依賴項¶
由於 ONNX 匯出器使用 onnx 和 onnxscript 將 PyTorch 運算子轉換為 ONNX 運算子，因此我們需要安裝它們。
pip install onnx pip install onnxscript
2. 建立一個簡單的圖像分類器模型¶
設定好環境後，讓我們開始使用 PyTorch 建立我們的圖像分類器模型，就像我們在 60 分鐘速成班中所做的一樣。
import torch
import torch.nn as nn
import torch.nn.functional as F
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
x = F.max_pool2d(F.relu(self.conv2(x)), 2)
x = torch.flatten(x, 1)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
3. 將模型匯出為 ONNX 格式¶
現在我們已經定義了我們的模型，我們需要實例化它並建立一個隨機的 32x32 輸入。 接下來，我們可以將模型匯出為 ONNX 格式。
torch_model = MyModel()
torch_input = torch.randn(1, 1, 32, 32)
onnx_program = torch.onnx.dynamo_export(torch_model, torch_input)
/usr/local/lib/python3.10/dist-packages/onnxscript/converter.py:823: FutureWarning:
'onnxscript.values.Op.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
/usr/local/lib/python3.10/dist-packages/onnxscript/converter.py:823: FutureWarning:
'onnxscript.values.OnnxFunction.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/_exporter_legacy.py:101: UserWarning:
torch.onnx.dynamo_export only implements opset version 18 for now. If you need to use a different opset version, please register them with register_custom_op.
/usr/local/lib/python3.10/dist-packages/torch/onnx/_internal/fx/onnxfunction_dispatcher.py:503: FutureWarning:
'onnxscript.values.TracedOnnxFunction.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
正如我們所看到的，我們不需要對模型進行任何程式碼變更。 產生的 ONNX 模型以二進位 protobuf 檔案的形式儲存在 torch.onnx.ONNXProgram 中。
4. 將 ONNX 模型儲存在檔案中¶
雖然在許多應用程式中將匯出的模型載入到記憶體中非常有用，但我們可以使用以下程式碼將其儲存到磁碟中
onnx_program.save("my_image_classifier.onnx")
您可以將 ONNX 檔案載入回記憶體，並使用以下程式碼檢查其格式是否正確
import onnx
onnx_model = onnx.load("my_image_classifier.onnx")
onnx.checker.check_model(onnx_model)
5. 使用 Netron 可視化 ONNX 模型圖¶
現在我們已經將模型儲存在檔案中，我們可以使用 Netron 對其進行可視化。 Netron 可以安裝在 macos、Linux 或 Windows 電腦上，也可以直接從瀏覽器執行。 讓我們透過開啟以下連結來嘗試 Web 版本：https://netron.app/。
Netron 開啟後，我們可以將我們的 my_image_classifier.onnx 檔案拖放到瀏覽器中，或者在點擊 開啟模型 按鈕後選擇它。
就是這樣！ 我們已成功將我們的 PyTorch 模型匯出為 ONNX 格式，並使用 Netron 對其進行可視化。
6. 使用 ONNX Runtime 執行 ONNX 模型¶
最後一步是使用 ONNX Runtime 執行 ONNX 模型，但在我們執行此操作之前，讓我們安裝 ONNX Runtime。
pip install onnxruntime
ONNX 標準不支援 PyTorch 支援的所有資料結構和類型，因此我們需要將 PyTorch 輸入調整為 ONNX 格式，然後再將其饋送到 ONNX Runtime。 在我們的範例中，輸入恰好是相同的，但在更複雜的模型中，它可能比原始 PyTorch 模型具有更多的輸入。
ONNX Runtime 需要一個額外的步驟，該步驟包含將所有 PyTorch 張量轉換為 Numpy (在 CPU 中)，並將它們包裝在一個字典中，該字典的鍵 (key) 是一個字串，字串內容為輸入名稱，值 (value) 則是 Numpy 張量。
現在我們可以建立一個ONNX Runtime 推論 Session，使用處理過的輸入執行 ONNX 模型，並取得輸出。在本教學中，ONNX Runtime 是在 CPU 上執行的，但它也可以在 GPU 上執行。
import onnxruntime
onnx_input = [torch_input]
print(f"Input length: {len(onnx_input)}")
print(f"Sample input: {onnx_input}")
ort_session = onnxruntime.InferenceSession("./my_image_classifier.onnx", providers=['CPUExecutionProvider'])
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
onnxruntime_input = {k.name: to_numpy(v) for k, v in zip(ort_session.get_inputs(), onnx_input)}
# onnxruntime returns a list of outputs
onnxruntime_outputs = ort_session.run(None, onnxruntime_input)[0]
Input length: 1
Sample input: [tensor([[[[-1.0416, 1.1125, -0.3602, ..., -0.0189, -1.2205, 0.3472],
[ 0.6651, 1.1037, -0.3674, ..., -1.4724, 0.2539, -0.0788],
[-0.1239, -0.6458, -0.7785, ..., -0.2674, 0.3019, -0.5682],
...,
[-0.0300, -0.4833, -0.3928, ..., -1.2406, 0.8488, -0.5473],
[-0.8185, -0.1276, 0.3475, ..., -1.0702, -1.6922, -0.6048],
[ 0.8268, -0.0248, -0.3354, ..., -0.9178, -0.3240, 0.7485]]]])]
7. 比較 PyTorch 結果與 ONNX Runtime 的結果¶
判斷導出的模型是否良好的最佳方式，是通過針對 PyTorch（我們的真值來源）進行數值評估。
為此，我們需要使用相同的輸入執行 PyTorch 模型，並將結果與 ONNX Runtime 的結果進行比較。在比較結果之前，我們需要轉換 PyTorch 的輸出，使其符合 ONNX 的格式。
torch_outputs = torch_model(torch_input)
assert len(torch_outputs) == len(onnxruntime_outputs)
for torch_output, onnxruntime_output in zip(torch_outputs, onnxruntime_outputs):
torch.testing.assert_close(torch_output, torch.tensor(onnxruntime_output))
print("PyTorch and ONNX Runtime output matched!")
print(f"Output length: {len(onnxruntime_outputs)}")
print(f"Sample output: {onnxruntime_outputs}")
PyTorch and ONNX Runtime output matched!
Output length: 1
Sample output: [[ 0.14531787 -0.05903321 -0.00652155 0.09054166 0.01458297 -0.08046442
-0.12109031 -0.03938238 -0.01814789 -0.01363543]]
結論¶
就是這樣！我們已成功將 PyTorch 模型導出為 ONNX 格式，將模型儲存到磁碟，使用 Netron 檢視它，使用 ONNX Runtime 執行它，並最終將其數值結果與 PyTorch 的數值結果進行比較。
延伸閱讀¶
下面的清單包含從基本範例到進階情境的教學，不一定按順序排列。您可以隨時直接跳到您感興趣的特定主題，或者耐心讀完所有內容，以了解關於 ONNX 匯出器的所有知識。
腳本總執行時間： (0 分鐘 1.749 秒)
