模型部署和推理架構學習筆記

一. 初步認識模型部署

1. 什么是ONNX？?

ONNX 就是一個?中間人?或?通用翻譯器。它讓你在喜歡的框架（如 PyTorch）里訓練好模型后，能輕松地把它變成一種?標準格式。然后，這個標準格式的模型可以被?很多不同的工具和硬件?(通過 ONNX Runtime 或其他支持 ONNX 的引擎) 理解和高效地運行，大大簡化了模型從訓練到實際應用的部署過程。它的目標是實現?“一次訓練，隨處運行”。

關鍵組成部分：

ONNX 格式：?基于 Protobuf 定義的文件格式，存儲了模型的網絡結構（計算圖）、參數（權重、偏置）和元數據。
ONNX 操作符集：?定義了一組標準的、不斷擴充的原子操作（如卷積、矩陣乘、激活函數等），這些操作是構建模型的基本單元。不同的框架在導出時，需要將其特有的操作映射到這些標準操作上。
ONNX 運行時：
- ONNX Runtime (ORT)：?由微軟維護的高性能推理引擎，專門用于在各種硬件平臺（CPU, GPU, FPGA, NPU等）上高效運行 ONNX 模型。它是 ONNX 生態中非常重要和流行的運行時。
- 其他運行時：?許多其他推理引擎和硬件加速庫也原生支持加載和運行 ONNX 模型。

主要優勢：

框架無關性：?打破訓練框架的壁壘。
硬件靈活性：?方便地將模型部署到多樣化的硬件環境。
部署效率：?避免了為每個目標平臺重復開發模型推理代碼。
生態系統：?擁有龐大的社區和眾多廠商支持（微軟、Meta、NVIDIA、Intel、AMD、Arm、華為、高通等）。
優化機會：?ONNX Runtime 等引擎可以對模型圖進行各種優化（如圖優化、算子融合、量化），顯著提升推理速度和降低資源消耗。

2. 模型部署

參考?https://zhuanlan.zhihu.com/p/516920606

在軟件工程中，模型部署是把開發完成的軟件投入使用的過程，包括環境配置、軟件安裝等步驟。那么對于深度學習來說，模型部署就是讓訓練好的模型在特定環境中運行的過程。會遇到的一些難題：

1）運行模型所需要的環境難以配置。深度學習模型通常是有一些框架編寫，比如PyTorch、Tensor Flow。由于框架規模、依賴環境的限制，這些框架不適合在手機、開發板等生產環境中安裝。

2）深度學習模型的結構通常比較龐大，需要大量的算力才能滿足實時運行的需求。模型運行效率需要優化。

因為這些難題，模型部署不能靠簡單的環境配置和安裝完成。流水線大致如下：

?為了讓模型最終能夠部署到某個環境上，開發者可以使用任意一種深度學習框架來定義網絡結構，并通過訓練確定網絡中的參數。之后，模型的結構和參數會被轉換成一種只描述網絡結構的中間表示，一些針對網絡結構的優化會在中間表示上進行。最后，用面向硬件的高性能編程框架（如CUDA，OpenCL）編寫，能高效執行深度學習網絡中算子的推理引擎會把中間表示轉換成特定的文件格式，并在對應的硬件平臺上高效運行模型。

這一條流水線解決了模型部署中兩大問題：使用對接深度學習框架和推理引擎的中間表示，開發者不必擔心如何在新環境中運行各個復雜的框架；通過中間表示的網絡結構優化和推理引擎對運算的底層優化，模型的運算效率大幅提升。

2.1 配置部署環境

讓我們用 PyTorch 實現一個超分辨率模型，并把模型部署到 ONNX Runtime 這個推理引擎上。

import os
import cv2
import numpy as np
import requests 
import torch 
import torch.onnx
from torch import nn

class SuperResolutionNet(nn.Module):def __init__(self, upscale_factor):super().__init__()self.upscale_factor = upscale_factorself.img_upsampler = nn.Upsample(scale_factor=self.upscale_factor,mode='bicubic',align_corners=False)self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4)self.conv2 = nn.Conv2d(64,32,kernel_size=1, padding=0)self.conv3 = nn.Conv2d(32,3,kernel_size=5,padding=2)self.relu = nn.ReLU()def forward(self, x):x = self.img_upsampler(x)out = self.relu(self.conv1(x))out = self.relu(self.conv2(out))out = self.conv3(out)return out
urls = ['https://download.openmmlab.com/mmediting/restorers/srcnn/srcnn_x4k915_1x16_1000k_div2k_20200608-4186f232.pth', 'https://raw.githubusercontent.com/open-mmlab/mmediting/master/tests/data/face/000001.png'] 
names = ['srcnn.pth', 'face.png']
for url,name in zip(urls, names):if not os.path.exists(name):open(name, 'wb').write(requests.get(url).content)def init_torch_model():torch_model = SuperResolutionNet(upscale_factor=3)state_dict = torch.load('srcnn.pth')['state_dict']for old_key in list(state_dict.keys()):new_key = '.'.join(old_key.split('.')[1:])state_dict[new_key] = state_dict.pop(old_key)torch_model.load_state_dict(state_dict)torch_model.eval()return torch_modelmodel = init_torch_model()
input_img = cv2.imread('face.png').astype(np.float32)input_img = np.transpose(input_img, [2, 0, 1])
input_img = np.expand_dims(input_img, 0)
torch_output = model(torch.from_numpy(input_img)).detach().numpy()# NCHW to HWC 
torch_output = np.squeeze(torch_output, 0) 
torch_output = np.clip(torch_output, 0, 255) 
torch_output = np.transpose(torch_output, [1, 2, 0]).astype(np.uint8) # Show image 
cv2.imwrite("face_torch.png", torch_output)

SRCNN 先把圖像上采樣到對應分辨率，再用 3 個卷積層處理圖像。為了方便起見，我們跳過訓練網絡的步驟，直接下載模型權重（由于 MMEditing 中 SRCNN 的權重結構和我們定義的模型不太一樣，我們修改了權重字典的 key 來適配我們定義的模型），同時下載好輸入圖片。為了讓模型輸出成正確的圖片格式，我們把模型的輸出轉換成 HWC 格式，并保證每一通道的顏色值都在 0~255 之間。如果腳本正常運行的話，一幅超分辨率的人臉照片會保存在 “face_torch.png” 中。

?在PyTorch模型測試正確后，我們來正式開始部署這個模型。我們下一步的任務是把pyTorch模型轉換成用中間表示ONNX描述的模型。

2.2 中間表示- ONNX

介紹ONNX之前，我們先從本質上來認識一下神經網絡的結構。神經網絡實際上知識描述了數據計算的過程，其結構可以用計算圖來表示。比如a + b可以用下圖的計算圖來表示：

為了加速計算，一些框架會使用對神經網絡“先編譯，后執行”的靜態圖來描述網絡。靜態圖的缺點是難以描述控制流（比如 if-else 分支語句和 for 循環語句），直接對其引入控制語句會導致產生不同的計算圖。比如循環執行n次a = a+b,對于不同的n,會產生不同的計算圖：

ONNX（Open Neural Network Exchange）是Facebook和微軟在2017年共同發布的，用于標準描述計算圖的一種格式。目前，在數家機構的共同維護下，ONNX已經對接了多種深度學習架構和多推理引擎。因此，ONNX被當成了深度學習框架到推理引擎的橋梁，就像編譯器的中間語言一樣。由于各個框架兼容性不一，我們通常只用ONNX表示更容易部署的靜態圖。

我們用下面的代碼把PyTorch的模型轉換成ONNX格式的模型：

x = torch.randn(1, 3, 256, 256) with torch.no_grad(): torch.onnx.export( model, x, "srcnn.onnx", opset_version=11, input_names=['input'], output_names=['output'])

其中，torch.onnx.export是PyTorch自帶的把模型轉換成ONNX格式的函數。讓我們先來看一下前三個必選參數：前三個參數分別是要轉換的模型、模型的任意一組輸入、導出的ONNX文件的文件名。轉換模型時，需要原模型和輸出文件迷宮是很容易理解的，但為什么需要為模型提供一組輸入呢？這就涉及到ONNX的轉換原理來。從PyTorch的模型到ONNX的模型，本質上是一種語言上的翻譯。直覺上的想法是像編譯器一樣徹底解析原模型的代碼，記錄所有控制流。但是前面我們通常只用ONNX記錄不考慮控制流的靜態圖。因此，PyTorch提供了一種叫做追蹤（trace）的模型轉換方法：給定一組輸入，再實際執行一遍模型，即把這組輸入對應的計算圖記錄下來，保存為ONNX格式。export函數用的就是追蹤導出方法，需要給任意一組輸入，讓模型跑起來。我們測試圖片的三通道，256*256大小的，這里也構造一個同樣形狀的隨機張量。?

剩下的參數中，opset_version表示ONNX算子集的版本。深度學習的發展會不斷的誕生新算子，為了支持這些新增的算子，ONNX會經常發布新的算子集。我們領opset_version=11，即使用第11個ONNX算子集，是因為SRC NN中的bicubic（雙三次插值）在opset11中才得到支持。剩下的兩個參數input_names, output_names是輸入、輸出tensor的名稱，我們稍后會用到這些名稱。

如果上述代碼運行成功，目錄下會新增一個"srcnn.onnx"的 ONNX 模型文件。我們可以用下面的腳本來驗證一下模型文件是否正確。

import onnx onnx_model = onnx.load("srcnn.onnx") 
try: onnx.checker.check_model(onnx_model) 
except Exception: print("Model incorrect") 
else: print("Model correct")

?其中，onnx.load函數用于讀取一個ONNX模型。onnx.checker.check_model用于檢查模型格式是否正確，如果有錯誤的話會直接報錯。我們的模型是正確的，控制臺中應該會打印出“Model correct”

接下來，讓我們來看一看 ONNX 模型具體的結構是怎么樣的。我們可以使用?Netron?（開源的模型可視化工具）來可視化 ONNX 模型。把 srcnn.onnx 文件從本地的文件系統拖入網站，即可看到如下的可視化結果：

點擊 input 或者 output，可以查看 ONNX 模型的基本信息，包括模型的版本信息，以及模型輸入、輸出的名稱和數據類型。

點擊某一個算子節點，可以看到算子的具體信息。比如點擊第一個 Conv 可以看到：

?每個算子記錄了算子屬性、圖結構、權重三類信息。

算子屬性信息即圖中attributs里的信息，對于卷積來說，算子屬性包括了卷積大小（kernel_shape）、卷積步長（strides）等內容。這些算子屬性最終會用來生成一個具體的算子。
圖結構信息指算子節點在就按圖中的名稱、鄰邊信息。對于圖中的卷積來說，該算子節點叫Conv_2,輸入數據叫做11，輸出數據叫做12。根據每個算子節點的圖結構信息，就能完整地復原出網絡的計算圖。
權重信息指的是網絡經過訓練后，算子存儲的權重信息。對于卷積來說，權重信息包括卷積核的權重值和卷積后的偏差值。點擊圖中 conv1.weight, conv1.bias 后面的加號即可看到權重信息的具體內容。

現在，我們有了SRCNN的ONNX模型。讓我們看看最后該如何把這個模型運行起來。

2.3 推理引擎- ONNX Runtime

ONNX Runtime是由微軟維護的一個跨平臺機器學習推理加速器，也就是我們前面提到的“推理引擎”。ONNX Runtime是直接對接ONNX的，即ONNX Runtime是直接對接ONNX的，即ONNX Runtime可以直接讀取并運行.onnx文件，而不需要再把.onnx格式的文件轉換成其他格式的文件。也就是說，對于PyTorch-ONNX-ONNX Runtime這條部署流程線，只要在目標設備中得到.onnx文件，并在ONNX Runtime上運行模型，模型部署就算大功告成了。

通過剛剛的操作，我們把PyTorch編寫的模型轉換成ONNX，并通過可視化檢查了模型的正確性。最后讓我們用ONNX Runtime運行一下模型，完成模型部署的最后一步。

ONNX Runtime提供了Python接口。接著剛才的腳本，我們可以添加如下代碼運行模型.

import onnxruntime ort_session = onnxruntime.InferenceSession("srcnn.onnx") 
ort_inputs = {'input': input_img} 
ort_output = ort_session.run(['output'], ort_inputs)[0] ort_output = np.squeeze(ort_output, 0) 
ort_output = np.clip(ort_output, 0, 255) 
ort_output = np.transpose(ort_output, [1, 2, 0]).astype(np.uint8) 
cv2.imwrite("face_ort.png", ort_output)

這段代碼中，出去后處理操作外，和ONNX Runtime相關的代碼只有三行。讓我們簡單解析一下這三行代碼。onnxruntime.InferenceSession用于獲取一個ONNX Runtime推理器，其參數是用于推理的ONNX模型文件。推理器的run方法用于模型推理，其第一個參數為輸出張量名的列表，第二個參數為輸入值的字典。其中輸入值字典的key為張量名，value為numpy類型的張量值。輸入輸出張量的名稱需要和torch.onnx.export中設置的輸入輸出名對應。

如果代碼正常運行的話，另一幅超分辨率照片會保存在"face_ort.png"中。這幅圖片和剛剛得到的"face_torch.png"是一模一樣的。這說明 ONNX Runtime 成功運行了 SRCNN 模型，模型部署完成了！以后有用戶想實現超分辨率的操作，我們只需要提供一個 "srcnn.onnx" 文件，并幫助用戶配置好 ONNX Runtime 的 Python 環境，用幾行代碼就可以運行模型了。或者還有更簡便的方法，我們可以利用 ONNX Runtime 編譯出一個可以直接執行模型的應用程序。我們只需要給用戶提供 ONNX 模型文件，并讓用戶在應用程序選擇要執行的 ONNX 模型文件名就可以運行模型了。

總結

模型部署，指把訓練好的模型在特定的環境中運行的結果。模型部署要解決模型框架兼容性差和模型運行速度慢兩大問題。
模型部署的常見流水線是“深度學習框架-中間表示-推理引擎”。其中比較常見的一個中間表示是ONNX。
深度學習模型實際上就是一個計算圖。模型部署時通常把模型轉換成靜態的就按圖，即沒有控制流（分支語句和循環語句）的計算圖。
PyTorch框架自帶對ONNX的支持，只需要構造一組隨機的輸入，并對模型調用torch.onnx.export即可完成Pytorch到ONNX的轉換。
推理引擎ONNX Runtime對ONNX模型有原生的支持。給定一個.onnx文件，只需要簡單的使用ONNX Runtime的Python API就可以完成模型推理。

二. 模型部署中常見的難題

一般模型部署會碰到以下幾類困難：

模型的動態化。出于性能的考慮，各推理框架都默認模型的輸入形狀、輸出形狀、結構時靜態的。而為了讓模型的泛用性更強，部署時需要在盡可能不影響原來邏輯的前提下，讓模型的輸入輸出或者結構動態話。
新算子的實現。深度學習日新月異，提出新算子的速度往往快于ONNX維護者支持的速度。為了部署新模型，部署工程師往往需要自己在ONNX和推理引擎中支持新算子。
中間表示于推理引擎的兼容問題。由于各推理引擎的實現不同，對ONNX難以形成統一的支持。為了確保模型在不同的推理引擎中有同樣的運行效果，部署工程師往往得為某個推理引擎定制模型代碼，這為模型部署引入了許多工作量。

3.1 實現動態放大的超分辨率模型

在原來的SRCNN中，圖片的放大比例是寫死在模型里的

def init_torch_model(): torch_model = SuperResolutionNet(upscale_factor=3)

我們使用 upscale_factor 來控制模型的放大比例。初始化模型的時候，我們默認令 upscale_factor 為 3，生成了一個放大 3 倍的 PyTorch 模型。這個 PyTorch 模型最終被轉換成了 ONNX 格式的模型。如果我們需要一個放大 4 倍的模型，需要重新生成一遍模型，再做一次到 ONNX 的轉換。

現在我們希望圖片的放大倍數可以自由設置。而我們給用戶的只有一個.onnx文件和運行超分辨率模型的應用程序。我們不修改.onnx文件的前提下改變放大倍數。

因此，我們必須修改原來的模型，令模型的放大倍數變成推理時的輸入。

import os
import cv2
import numpy as np
import requests 
import torch 
import torch.onnx
from torch import nn
from torch.nn.functional import interpolateclass SuperResolutionNet(nn.Module):def __init__(self, upscale_factor):super().__init__()self.upscale_factor = upscale_factorself.img_upsampler = nn.Upsample(scale_factor=self.upscale_factor,mode='bicubic',align_corners=False)self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4)self.conv2 = nn.Conv2d(64,32,kernel_size=1, padding=0)self.conv3 = nn.Conv2d(32,3,kernel_size=5,padding=2)self.relu = nn.ReLU()def forward(self, x, upscale_factor):x = interpolate(x,scale_factor = upscale_factor,mode ='bicubic',align_corners=False)out = self.relu(self.conv1(x))out = self.relu(self.conv2(out))out = self.conv3(out)return out
urls = ['https://download.openmmlab.com/mmediting/restorers/srcnn/srcnn_x4k915_1x16_1000k_div2k_20200608-4186f232.pth', 'https://raw.githubusercontent.com/open-mmlab/mmediting/master/tests/data/face/000001.png'] 
names = ['srcnn.pth', 'face.png']
for url,name in zip(urls, names):if not os.path.exists(name):open(name, 'wb').write(requests.get(url).content)def init_torch_model():torch_model = SuperResolutionNet(upscale_factor=3)state_dict = torch.load('srcnn.pth')['state_dict']for old_key in list(state_dict.keys()):new_key = '.'.join(old_key.split('.')[1:])state_dict[new_key] = state_dict.pop(old_key)torch_model.load_state_dict(state_dict)torch_model.eval()return torch_modelmodel = init_torch_model()
input_img = cv2.imread('face.png').astype(np.float32)input_img = np.transpose(input_img, [2, 0, 1])
input_img = np.expand_dims(input_img, 0)
torch_output = model(torch.from_numpy(input_img), 3).detach().numpy()# NCHW to HWC 
torch_output = np.squeeze(torch_output, 0) 
torch_output = np.clip(torch_output, 0, 255) 
torch_output = np.transpose(torch_output, [1, 2, 0]).astype(np.uint8) # Show image 
cv2.imwrite("face_torch.png", torch_output)

SuperResolutionNet 未修改之前，nn.Upsample 在初始化階段固化了放大倍數，而 PyTorch 的 interpolate 插值算子可以在運行階段選擇放大倍數。因此，我們在新腳本中使用 interpolate 代替 nn.Upsample，從而讓模型支持動態放大倍數的超分。在使用模型推理時，我們把放大倍數設置為 3。

torch_output = model(torch.from_numpy(input_img), 3).detach().numpy()

最后，圖片保存在文件 "face_torch_2.png" 中。一切正常的話，"face_torch_2.png" 和 "face_torch.png" 的內容一模一樣。?

導出模型時：

x = torch.randn(1, 3, 256, 256)with torch.no_grad():torch.onnx.export(model, (x,3),"srcnn2.onnx",opset_version=11,input_names=['input', 'factor'],output_names = ['output'])

運行這些腳本時，會報一長串錯誤。沒辦法，我們碰到了模型部署中的兼容性問題。

3.2 自定義算子

直接使用Pytorch模型的話，我們修改幾行代碼就能實現模型輸入的動態化。但在模型部署中，我們要花數倍的時間來設法解決這一問題。現在，讓我們順著解決問題的思路，體驗一下模型部署的困難，并學習使用自定義算子的方式，解決超分辨率模型的動態化問題。

剛剛的報錯是因為 PyTorch 模型在導出到 ONNX 模型時，模型的輸入參數的類型必須全部是 torch.Tensor。而實際上我們傳入的第二個參數" 3 "是一個整形變量。這不符合 PyTorch 轉 ONNX 的規定。我們必須要修改一下原來的模型的輸入。為了保證輸入的所有參數都是 torch.Tensor 類型的，我們做如下修改：

... class SuperResolutionNet(nn.Module): def forward(self, x, upscale_factor): x = interpolate(x, scale_factor=upscale_factor.item(), mode='bicubic', align_corners=False) ... # Inference 
# Note that the second input is torch.tensor(3) 
torch_output = model(torch.from_numpy(input_img), torch.tensor(3)).detach().numpy() ... with torch.no_grad(): torch.onnx.export(model, (x, torch.tensor(3)), "srcnn2.onnx", opset_version=11, input_names=['input', 'factor'], output_names=['output'])

由于 PyTorch 中 interpolate 的 scale_factor 參數必須是一個數值，我們使用 torch.Tensor.item() 來把只有一個元素的 torch.Tensor 轉換成數值。之后，在模型推理時，我們使用 torch.tensor(3) 代替 3，以使得我們的所有輸入都滿足要求。現在運行腳本的話，無論是直接運行模型，還是導出 ONNX 模型，都不會報錯了。?

但是，導出 ONNX 時卻報了一條 TraceWarning 的警告。這條警告說有一些量可能會追蹤失敗。

/var/folders/pd/9txrcrys3rdfqk4hyxxtvszr0000gn/T/ipykernel_28173/3437727912.py:16: TracerWarning: Converting a tensor to a Python number might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs! scale_factor = upscale_factor.item(),

這是怎么回事呢？讓我們把生成的 srcnn2.onnx 用?Netron?可視化一下：

可以發現，雖然我們把模型推理的輸入設置為兩個，但是ONNX模型還是長的和原來一摸一樣，很自由一個叫input的輸入，這是由于我們使用了torch.Tensor.item()把數據從Tensor里取出來，而導出ONNX模型時這個操作時無法被記錄的，只好報了一條Trace Warning。這導致interpolate插值函數的放大倍數還是被設置成了“3”這個固定值，我們導出的“srcnn2.onnx”和最開始的“srcnn.onnx”完全相同。

直接修改原來的模型似乎行不通，我們得從 PyTorch 轉 ONNX 的原理入手，強行令 ONNX 模型明白我們的想法了。?

仔細觀察 Netron 上可視化出的 ONNX 模型，可以發現在 PyTorch 中無論是使用最早的 nn.Upsample，還是后來的 interpolate，PyTorch 里的插值操作最后都會轉換成 ONNX 定義的 Resize 操作。也就是說，所謂的PyTorch轉ONNX實際上就是把每個PyTorch的操作映射成了ONNX定義的算子。

其中，展開scales，可以看到scales是一個長度為4的一維張量，其內容為[1,1,3,3]

一維張量[1, 1, 3, 3]表示Resize操作每一個維度的縮放系數；其類型為Initializer，表示這個值是根據常量初始化出來的。如果我們能夠自己生成一個ONNX的Resize算子，讓scales成為一個可變量而不是常量，就像它上面的X一樣，那這個超分辨模型就能動態縮放了。

現有實現插值的 PyTorch 算子有一套規定好的映射到 ONNX Resize 算子的方法，這些映射出的 Resize 算子的 scales 只能是常量，無法滿足我們的需求。我們得自己定義一個實現插值的 PyTorch 算子，然后讓它映射到一個我們期望的 ONNX Resize 算子上。?

下面的腳本定義了一個 PyTorch 插值算子，并在模型里使用了它。我們先通過運行模型來驗證該算子的正確性：

class NewInterpolate(torch.autograd.Function): @staticmethod def symbolic(g, input, scales): return g.op("Resize", input, g.op("Constant", value_t=torch.tensor([], dtype=torch.float32)), scales, coordinate_transformation_mode_s="pytorch_half_pixel", cubic_coeff_a_f=-0.75, mode_s='cubic', nearest_mode_s="floor") @staticmethod def forward(ctx, input, scales): scales = scales.tolist()[-2:] return interpolate(input, scale_factor=scales, mode='bicubic',

def symbolic(g, input, scales):

g：表示一個圖（graph）上下文，用于構建ONNX圖
input：要調整大小的輸入張量
scales：縮放因子，是一個一維張量（長度為輸入維度數），指定每個維度上的縮放比例

參數詳解：

"Resize"：要創建的ONNX算子的類型（調整大小操作）
input：輸入張量，即需要調整大小的數據
g.op("Constant", value_t=torch.tensor([], dtype=torch.float32))：
????????創建一個空的常量張量作為roi（感興趣區域）參數
????????這里使用空張量表示使用整個輸入區域
????????dtype=torch.float32指定數據類型為32位浮點數
scales：
1. 縮放因子張量
  例如：對于圖像數據（N, C, H, W），scales可能為[1.0, 1.0, 2.0, 2.0]表示高度和寬度各放大2倍
coordinate_transformation_mode_s="pytorch_half_pixel"：
- 坐標變換模式
  pytorch_half_pixel：使用PyTorch風格的半像素坐標變換
  確保ONNX的resize操作與PyTorch的interpolate()行為一致
cubic_coeff_a_f=-0.75：
- 三次插值的系數
  -0.75是常用的值（對應Catmull-Rom樣條插值）
  僅當插值模式為'cubic'時生效
mode_s='cubic'：
- 插值模式
  'cubic'表示使用三次樣條插值
  其他可能值：'nearest', 'linear'等
nearest_mode_s="floor"：
- 當插值模式為nearest時使用的舍入方法
  "floor"表示向下取整
  這里雖然模式是cubic，但ONNX要求必須指定此參數

先理清一下思路，我們希望新的插值算子有兩個輸入，一個是被用于操作的圖像，一個是圖像的放縮比例。前面講到，為了對接ONNX中的Resize算子的scales參數，這個放縮比例是一個【1, 1, x, x】張量，其中x為放大倍數。在之前放大3倍的模型中，這個參數被固定成了【1，1，3，3】。因此，在插值算子中，我們希望模型的第二個輸入是一個【1，1，w，h】的張量，其中w和h分別為圖片寬和高的放大倍數。

搞清楚了插值算子的輸入，再看一看算子的具體實現。算子的推理行為由算子的forward方法決定。該方法的第一個參數必須為ctx,后面的參數為算子的自定義輸入，我們設置兩個輸入，分別為被操作的圖像和放縮比例。為保證推理正確，需要把【1，1，w, h】格式的輸入對接到原來的interpolate函數上。我們的做法是截取輸入張量的后兩個元素，把這兩個元素以list的格式傳入interpolate的scale_factor參數。

接下來，我們要決定新算子映射到ONNX算子的方法，映射到ONNX的方法由一個算子的symbolic方法決定。symbolic方法的第一個參數必須是g,之后的參數是算子的自定義輸入，和forward函數一樣。ONNX算子的具體定義由g.op實現。g.op的每個參數都可以映射到ONNX中的算子屬性：

?對于其他參數，我們可以照著現在的Resize算子填。而要注意的是，我們現在希望scales參數是由輸入動態決定的。因此，在填入ONNX的scales時，我們要把symbolic方法的輸入參數中scales填入。

接著，讓我們把新模型導出成ONNX模型：

x = torch.randn(1, 3, 256, 256)
with torch.no_grad():torch.onnx.export(model, (x, factor),"srcnn3.onnx",opset_version=11,input_names=['input', 'factor'],output_names=['output'])

可以看到，正如我們所期望的，導出的 ONNX 模型有了兩個輸入！第二個輸入表示圖像的放縮比例。

運行上面的代碼，可以得到一個邊長放大4倍的超分辨率圖片 "face_ort_3.png"。動態的超分辨率模型生成成功了！只要修改 input_factor，我們就可以自由地控制圖片的縮放比例。?

我們剛剛的工作，實際上是繞過 PyTorch 本身的限制，憑空“捏”出了一個 ONNX 算子。事實上，我們不僅可以創建現有的 ONNX 算子，還可以定義新的 ONNX 算子以拓展 ONNX 的表達能力。后續教程中我們將介紹自定義新 ONNX 算子的方法。

?總結

模型部署中常見的幾類困難有：模型的動態化；新算子的實現；框架間的兼容。
PyTorch轉ONNX，實際上就是把每一個操作轉化成ONNX定義的某一個算子。比如對于PyTorch中的Upsample和interpolate,在轉ONNX后最終都會成為ONNX的Resize算子。
通過修改繼承自torch.autograd.Function算子的symbolic方法，可以改變該算子映射到ONNX算子的行為。

三. PyTorch轉ONNX詳解

ONNX 是目前模型部署中最重要的中間表示之一。學懂了 ONNX 的技術細節，就能規避大量的模型部署問題。?
在把 PyTorch 模型轉換成 ONNX 模型時，我們往往只需要輕松地調用一句torch.onnx.export就行了。這個函數的接口看上去簡單，但它在使用上還有著諸多的“潛規則”。在這篇教程中，我們會詳細介紹 PyTorch 模型轉 ONNX 模型的原理及注意事項。除此之外，我們還會介紹 PyTorch 與 ONNX 的算子對應關系，以教會大家如何處理 PyTorch 模型轉換時可能會遇到的算子支持問題。

4.1 torch.onnx.export細節

在這一節里，我們將詳細介紹 PyTorch 到 ONNX 的轉換函數——?torch.onnx.export。我們希望大家能夠更加靈活地使用這個模型轉換接口，并通過了解它的實現原理來更好地應對該函數的報錯（由于模型部署的兼容性問題，部署復雜模型時該函數時常會報錯）。

1. 計算圖導出方法

（幫助理解為主，新的pytorch計算圖有所改變）

TorchScript 是一種序列化和優化PyTorch模型的格式，在優化過程中，一個torch.nn.Module模型會被轉換成TorchScript的torch.jit.ScriptModule模型。現在，TorchScript也被當成一種中間表示使用。

torch.onnx.export中需要的模型實際上是一個torch.jit.ScriptModule。而把普通的PyTorch模型轉成一個這樣的TorchScript模型，有跟蹤（trace）和記錄（script）兩種導出計算圖的方法。如果給torch.onnx.export傳入了一個普通PyTorch模型（torch.nn.Module）,那么這個模型會默認使用trace方法導出：

回憶一下我們第一篇教程知識：跟蹤法只能通過實際運行一遍模型的方法導出模型的靜態圖，即無法識別出模型中的控制流（如循環）；記錄法則能通過解析模型來正確記錄所有的控制流。我們以下面這段代碼為例來看一看這兩種轉換方法的區別：

import torch class Model(torch.nn.Module): def __init__(self, n): super().__init__() self.n = n self.conv = torch.nn.Conv2d(3, 3, 3) def forward(self, x): for i in range(self.n): x = self.conv(x) return x models = [Model(2), Model(3)] 
model_names = ['model_2', 'model_3'] for model, model_name in zip(models, model_names): dummy_input = torch.rand(1, 3, 10, 10) dummy_output = model(dummy_input) model_trace = torch.jit.trace(model, dummy_input) model_script = torch.jit.script(model) torch.onnx.export(model_trace, dummy_input, f'{model_name}_trace.onnx',input_names=['input'],        # 新增：指定輸入名稱output_names=['output'],      # 新增：指定輸出名稱dynamic_axes={                # 新增：處理動態維度'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})torch.onnx.export(model_script, dummy_input, f'{model_name}_script.onnx',input_names=['input'],        # 新增：指定輸入名稱output_names=['output'],      # 新增：指定輸出名稱dynamic_axes={                # 新增：處理動態維度'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})

在這段代碼里，我們定義了一個帶循環的模型，模型通過參數n來控制輸入張量被卷積的次數。之后，我們各創建了一個n=2和n=3的模型。我們把這兩個模型分別用跟蹤和記錄的方法進行導出。

值得一提的是，由于這里的兩個模型（model_trace, model_script）是TorchScript模型，export函數已經不需要再運行一遍模型了。（如果模型是用跟蹤法得到的，那么在執行torch.jit.trace的時候就運行過一遍了；而用記錄法導出時，模型不需要實際運行）參數中的dummy input和dummy output僅僅是為了獲取輸入和輸出張量的類型和形狀。

運行上面的代碼，我們把得到的4個onnx文件用Netron可視化：

首先看跟蹤法得到的ONNX模型結構，可以看出來，對于不同的n,ONNX模型結構式不一樣的。?

?而用記錄法的話，最終的ONNX模型用Loop節點來表示循環。這樣哪怕對于不同的n，ONNX模型也有同樣的結構。

?2. 參數講解

了解完轉換函數的原理后，我們來詳細介紹一下該函數的主要參數作用：

function) def export(model: Module | ExportedProgram | ScriptModule | ScriptFunction,args: tuple[Any, ...] = (),f: str | PathLike | None = None,*,kwargs: dict[str, Any] | None = None,export_params: bool = True,verbose: bool | None = None,input_names: Sequence[str] | None = None,output_names: Sequence[str] | None = None,opset_version: int | None = None,dynamic_axes: Mapping[str, Mapping[int, str]] | Mapping[str, Sequence[int]] | None = None,keep_initializers_as_inputs: bool = False,dynamo: bool = False,
external_data: bool = True,dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any] | None = None,custom_translation_table: dict[(...) -> Any, ((...) -> Any) | Sequence[(...) -> Any]] | None = None,report: bool = False,optimize: bool = True,verify: bool = False,profile: bool = False,dump_exported_program: bool = False,artifacts_dir: str | PathLike = ".",fallback: bool = False,training: TrainingMode = _C_onnx.TrainingMode.EVAL,operator_export_type: OperatorExportTypes = _C_onnx.OperatorExportTypes.ONNX,do_constant_folding: bool = True,custom_opsets: Mapping[str, int] | None = None,export_modules_as_functions: bool | Collection[type[Module]] = False,autograd_inlining: bool = True
) -> (ONNXProgram | None)

input_names, output_names?

設置輸入和輸出張量的名稱。如果不設置的話，會自動分配一些簡單的名字（如數字）。?
ONNX 模型的每個輸入和輸出張量都有一個名字。很多推理引擎在運行 ONNX 文件時，都需要以“名稱-張量值”的數據對來輸入數據，并根據輸出張量的名稱來獲取輸出數據。在進行跟張量有關的設置（比如添加動態維度）時，也需要知道張量的名字。?
在實際的部署流水線中，我們都需要設置輸入和輸出張量的名稱，并保證 ONNX 和推理引擎中使用同一套名稱。

opset_version?

轉換時參考哪個 ONNX 算子集版本，默認為 9。后文會詳細介紹 PyTorch 與 ONNX 的算子對應關系。?

dynamic_axes?

指定輸入輸出張量的哪些維度是動態的。?
為了追求效率，ONNX 默認所有參與運算的張量都是靜態的（張量的形狀不發生改變）。但在實際應用中，我們又希望模型的輸入張量是動態的，尤其是本來就沒有形狀限制的全卷積模型。因此，我們需要顯式地指明輸入輸出張量的哪幾個維度的大小是可變的。?
我們來看一個dynamic_axes的設置例子：

import torch class Model(torch.nn.Module): def __init__(self): super().__init__() self.conv = torch.nn.Conv2d(3, 3, 3) def forward(self, x): x = self.conv(x) return x model = Model() 
dummy_input = torch.rand(1, 3, 10, 10) 
model_names = ['model_static.onnx',  
'model_dynamic_0.onnx',  
'model_dynamic_23.onnx'] dynamic_axes_0 = { 'in' : [0], 'out' : [0] 
} 
dynamic_axes_23 = { 'in' : [2, 3], 'out' : [2, 3] 
} torch.onnx.export(model, dummy_input, model_names[0],  
input_names=['in'], output_names=['out']) 
torch.onnx.export(model, dummy_input, model_names[1],  
input_names=['in'], output_names=['out'], dynamic_axes=dynamic_axes_0) 
torch.onnx.export(model, dummy_input, model_names[2],  
input_names=['in'], output_names=['out'], dynamic_axes=dynamic_axes_23)

首先，我們導出 3 個 ONNX 模型，分別為沒有動態維度、第 0 維動態、第 2 第 3 維動態的模型。?
在這份代碼里，我們是用列表的方式表示動態維度，例如：?

dynamic_axes_0 = { 'in' : [0], 'out' : [0] 
}

由于 ONNX 要求每個動態維度都有一個名字，這樣寫的話會引出一條 UserWarning，警告我們通過列表的方式設置動態維度的話系統會自動為它們分配名字。一種顯式添加動態維度名字的方法如下：?

dynamic_axes_0 = { 'in' : {0: 'batch'}, 'out' : {0: 'batch'} 
}

我們在模型導出計算圖時用的是一個形狀為(1, 3, 10, 10)的張量。現在，我們來嘗試以形狀分別是(1, 3, 10, 10), (2, 3, 10, 10), (1, 3, 20, 20)為輸入，用ONNX Runtime運行一下這幾個模型，看看哪些情況下會報錯，并保存對應的報錯信息。得到的輸出信息應該如下：?

Input[0] on model model_static.onnx succeed. 
Input[1] on model model_static.onnx error. 
Input[2] on model model_static.onnx error. 
Input[0] on model model_dynamic_0.onnx succeed. 
Input[1] on model model_dynamic_0.onnx succeed. 
Input[2] on model model_dynamic_0.onnx error. 
Input[0] on model model_dynamic_23.onnx succeed. 
Input[1] on model model_dynamic_23.onnx error. 
Input[2] on model model_dynamic_23.onnx succeed.

可以看出，形狀相同的(1, 3, 10, 10)的輸入在所有模型上都沒有出錯。而對于batch（第 0 維）或者長寬（第 2、3維）不同的輸入，只有在設置了對應的動態維度后才不會出錯。我們可以錯誤信息中找出是哪些維度出了問題。比如我們可以用以下代碼查看input[1]在model_static.onnx中的報錯信息：?

print(exceptions[(1, 'model_static.onnx')]) # output 
# [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: in for the following indices index: 0 Got: 2 Expected: 1 Please fix either the inputs or the model.

這段報錯告訴我們名字叫in的輸入的第 0 維不匹配。本來該維的長度應該為 1，但我們的輸入是 2。實際部署中，如果我們碰到了類似的報錯，就可以通過設置動態維度來解決問題。

完整的export例子：

import torch
import torch.onnx# 定義模型
class MyModel(torch.nn.Module):def forward(self, x):return torch.nn.functional.relu(x)model = MyModel().eval()# 導出配置
torch.onnx.export(model=model,args=torch.randn(1, 3, 224, 224),f="my_model.onnx",input_names=["input"],output_names=["output"],dynamic_axes={"input": {0: "batch_size"},"output": {0: "batch_size"}},opset_version=14,do_constant_folding=True,verbose=True,training=torch.onnx.TrainingMode.EVAL
)

?3. 使用提示

通過學習之前的知識，我們基本掌握了?torch.onnx.export函數的部分實現原理和參數設置方法，足以完成簡單模型的轉換了。但在實際應用中，使用該函數還會踩很多坑。這里我們模型部署團隊把在實戰中積累的一些經驗分享給大家。

使模型在 ONNX 轉換時有不同的行為?

有些時候，我們希望模型在導出至 ONNX 時有一些不同的行為模型在直接用 PyTorch 推理時有一套邏輯，而在導出的ONNX模型中有另一套邏輯。比如，我們可以把一些后處理的邏輯放在模型里，以簡化除運行模型之外的其他代碼。torch.onnx.is_in_onnx_export()可以實現這一任務，該函數僅在執行?torch.onnx.export()時為真。以下是一個例子：

import torch class Model(torch.nn.Module): def __init__(self): super().__init__() self.conv = torch.nn.Conv2d(3, 3, 3) def forward(self, x): x = self.conv(x) if torch.onnx.is_in_onnx_export(): x = torch.clip(x, 0, 1) return x

這里，我們僅僅在模型導出時把輸出張量的數值限制在【0，1】之間。使用is_in_onnx_Export確實能讓我們方便地在代碼中添加和模型部署相關的邏輯。但是，這些代碼對只關心模型訓練的開發者和用戶來說很不友好，突兀的部署邏輯會降低代碼整體的可讀性。同時，is_in_onnx_export只能在每個需要添加的部署邏輯的地方“打補丁”，難以進行統一的管理。（MMDeploy重新機制可以規避這些問題）。

4.2 PyTorch 對 ONNX 的算子支持

在確保torch.onnx.export()的調用方法無誤后，PyTorch 轉 ONNX 時最容易出現的問題就是算子不兼容了。這里我們會介紹如何判斷某個 PyTorch 算子在 ONNX 中是否兼容，以助大家在碰到報錯時能更好地把錯誤歸類。而具體添加算子的方法我們會在之后的文章里介紹。?
在轉換普通的torch.nn.Module模型時，PyTorch 一方面會用跟蹤法執行前向推理，把遇到的算子整合成計算圖；另一方面，PyTorch 還會把遇到的每個算子翻譯成 ONNX 中定義的算子。在這個翻譯過程中，可能會碰到以下情況：?

該算子可以一對一地翻譯成一個 ONNX 算子。?
該算子在 ONNX 中沒有直接對應的算子，會翻譯成一至多個 ONNX 算子。?
該算子沒有定義翻譯成 ONNX 的規則，報錯。?

那么，該如何查看 PyTorch 算子與 ONNX 算子的對應情況呢？由于 PyTorch 算子是向 ONNX 對齊的，這里我們先看一下 ONNX 算子的定義情況，再看一下 PyTorch 定義的算子映射關系。

ONNX算子文檔

ONNX 算子的定義情況，都可以在官方的算子文檔中查看。這份文檔十分重要，我們碰到任何和 ONNX 算子有關的問題都得來”請教“這份文檔。https://github.com/onnx/onnx/blob/main/docs/Operators.md

這份文檔中最重要的開頭的這個算子變更表格。表格的第一列是算子名，第二列是該算子發生變動的算子集版本號，也就是我們之前在torch.onnx.export中提到的opset_version表示的算子集版本號。通過查看算子第一次發生變動的版本號，我們可以知道某個算子是從哪個版本開始支持的；通過查看某算子小于等于opset_version的第一個改動記錄，我們可以知道當前算子集版本中該算子的定義規則。

通過點擊表格中的鏈接，我們可以查看某個算子的輸入、輸出參數規定及使用示例。比如上圖是 Relu 在 ONNX 中的定義規則，這份定義表明 Relu 應該有一個輸入和一個輸入，輸入輸出的類型相同，均為 tensor。

4.3 PyTorch 對 ONNX 算子的映射?

在 PyTorch 中，和 ONNX 有關的定義全部放在?torch.onnx目錄中

https://github.com/pytorch/pytorch/tree/main/torch/onnx

其中，symbolic_opset{n}.py（符號表文件）即表示 PyTorch 在支持第 n 版 ONNX 算子集時新加入的內容。我們之前講過， bicubic 插值是在第 11 個版本開始支持的。我們以它為例來看看如何查找算子的映射情況。?
首先，使用搜索功能，在torch/onnx文件夾搜索"bicubic"，可以發現這個這個插值在第 11 個版本的定義文件中：

我們按照代碼的調用邏輯，逐步跳轉直到最底層的 ONNX 映射函數：?

upsample_bicubic2d = _interpolate("upsample_bicubic2d", 4, "cubic") -> def _interpolate(name, dim, interpolate_mode): return sym_help._interpolate_helper(name, dim, interpolate_mode) -> def _interpolate_helper(name, dim, interpolate_mode): def symbolic_fn(g, input, output_size, *args): ... return symbolic_fn

最后，在symbolic_fn中，我們可以看到插值算子是怎么樣被映射成多個 ONNX 算子的。其中，每一個g.op就是一個 ONNX 的定義。比如其中的?Resize?算子就是這樣寫的：?

return g.op("Resize", input, empty_roi, empty_scales, output_size, coordinate_transformation_mode_s=coordinate_transformation_mode, cubic_coeff_a_f=-0.75,  # only valid when mode="cubic" mode_s=interpolate_mode,  # nearest, linear, or cubic nearest_mode_s="floor")  # only valid when mode="nearest"

通過在前面提到的ONNX 算子文檔中查找 Resize 算子的定義，我們就可以知道這每一個參數的含義了。用類似的方法，我們可以去查詢其他 ONNX 算子的參數含義，進而知道 PyTorch 中的參數是怎樣一步一步傳入到每個 ONNX 算子中的。?
掌握了如何查詢 PyTorch 映射到 ONNX 的關系后，我們在實際應用時就可以在?torch.onnx.export()的opset_version中先預設一個版本號，碰到了問題就去對應的 PyTorch 符號表文件里去查。如果某算子確實不存在，或者算子的映射關系不滿足我們的要求，我們就可能得用其他的算子繞過去，或者自定義算子了。

總結

跟蹤法和記錄法在導出帶控制語句的計算圖時有什么區別。?
torch.onnx.export()中該如何設置?input_names, output_names, dynamic_axes。?
使用?torch.onnx.is_in_onnx_export()來使模型在轉換到 ONNX 時有不同的行為。?
如何查詢 ONNX 算子文檔（https://github.com/onnx/onnx/blob/main/docs/Operators.md）。?
如何查詢 PyTorch 對某個 ONNX 版本的新特性支持情況。?
如何判斷 PyTorch 對某個 ONNX 算子是否支持，支持的方法是怎樣的。

四.?在 PyTorch 中支持更多 ONNX 算子

模型部署入門系列教程持續更新啦，在上一篇教程中，我們系統地學習了 PyTorch 轉 ONNX 的方法，可以發現 PyTorch 對 ONNX 的支持還不錯。但在實際的部署過程中，難免碰到模型無法用原生 PyTorch 算子表示的情況。這個時候，我們就得考慮擴充 PyTorch，即在 PyTorch 中支持更多 ONNX 算子。?

而要使 PyTorch 算子順利轉換到 ONNX ，我們需要保證以下三個環節都不出錯：?

算子在 PyTorch 中有實現?
有把該 PyTorch 算子映射成一個或多個 ONNX 算子的方法?
ONNX 有相應的算子?

可在實際部署中，這三部分的內容都可能有所缺失。其中最壞的情況是：我們定義了一個全新的算子，它不僅缺少 PyTorch 實現，還缺少 PyTorch 到 ONNX 的映射關系。但所謂車到山前必有路，對于這三個環節，我們也分別都有以下的添加支持的方法：?

PyTorch 算子?
- 組合現有算子?
- 添加 TorchScript 算子?
- 添加普通 C++ 拓展算子?
映射方法?
- 為 ATen 算子添加符號函數?
- 為 TorchScript 算子添加符號函數?
- 封裝成?torch.autograd.Function?并添加符號函數?
ONNX 算子?
- 使用現有 ONNX 算子?
- 定義新 ONNX 算子

4.1 支持ATen算子

實際的部署過程中，我們都有可能會碰到一個最簡單的算子缺失問題：算子在ATen中已經實現了，ONNX中也有相關算子的定義，但是相關算子映射成ONNX的規則沒有寫。在這種情況下，我們只需要為ATen算子補充描述映射規則的符號函數就行了。

ATen是PyTorch內置的C++張量計算庫，PyTorch算子在底層絕大多數計算都是用ATen實現的。

Asinh算子在ATen中有實現，卻缺少了映射到ONNX算子的符號函數。在這里，我們來嘗試為它補充符號函數，并導出一個包含這個算子的ONNX模型。

獲取ATen中算子接口定義

為了編寫符號函數，我們需要獲得asinh推理接口的輸入參數定義。這時，我們要去torch/_C/_VariableFunctions.pyi?和?torch/nn/functional.pyi?這兩個文件中搜索我們剛剛得到的這個算子名。這兩個文件是編譯PyTorch時本地自動生成的文件，里面包含了ATen算子的PyTorch調用接口。通過搜索，我們可以知道?asinh?在文件?torch/_C/_VariableFunctions.pyi?中，其接口定義為:

def asinh(input: Tensor, *, out: Optional[Tensor]=None) -> Tensor: ...

經過這些步驟，我們確認了缺失的算子名為asinh, 它是一個有實現的ATen算子。我們還記下了asinh的調用接口。接下來，我們要為它補充符號函數，使它轉換成ONNX模型時不再報錯。?

添加符號函數

到目前為止，我們已經多次接觸了定義PyTorch到ONNX映射規則的符號函數了。

符號函數，可以堪稱PyTorch算子類的一個靜態方法。在把PyTorch模型轉換成ONNX模型時，各個PyTorch算子的符號函數會被依次調用，以完成PyTorch算子到ONNX算子的轉換。
符號函數的定義一般如下：

def symbolic(g: torch._C.Graph, input_0: torch._C.Value, input_1: torch._C.Value, ...):

其中，torch._C.Graph和torch._C.Value都對應PyTorch的C++實現里的一些類。我們只需要知道第一個參數就固定叫g,它表示和計算圖相關的內容；后面的每個參數都表示算子的輸入，需要和算子的前向推理接口的輸入相同。對于ATen算子來收，他們的向前推理接口就是上述兩個.pyi文件里的函數接口。

g有一個方法op。在把Pytorch算子轉換成ONNX算子時，需要在符號函數中調用此方法來為最終的計算圖添加一個ONNX算子。其定義如下：

def op(name: str, input_0: torch._C.Value, input_1: torch._C.Value, ...)

?其中，第一個參數時算子名稱。如果該算子是普通的ONNX算子，只需要把它在ONNX官方文檔里的名稱填進去即可。

在最簡單的情況下，我們只要把 PyTorch 算子的輸入用g.op()一一對應到 ONNX 算子上即可，并把g.op()的返回值作為符號函數的返回值。在情況更復雜時，我們轉換一個 PyTorch 算子可能要新建若干個 ONNX 算子。

補充完了背景知識，讓我們回到?asinh?算子上，來為它編寫符號函數。我們先去翻閱一下 ONNX 算子文檔，學習一下我們在符號函數里的映射關系?g.op()?里應該怎么寫。Asinh?的文檔寫道：該算子有一個輸入?input，一個輸出?output，二者的類型都為張量。

到這里，我們已經完成了信息收集環節。我們在上一小節得知了?asinh?的推理接口定義，在這一小節里收集了 ONNX 算子?Asinh?的定義。現在，我們可以用代碼來補充這二者的映射關系了。在剛剛導出??asinh?算子的代碼中，我們添加以下內容：

import torchclass Model(torch.nn.Module):def __init__(self):super().__init__()def forward(self, x):return torch.asinh(x)#from torch.onnx.symbolic_registry import register_op
from torch.onnx import register_custom_op_symbolicdef asinh_symbolic(g, input, *, out=None):return g.op("Asinh", input)
# 注冊到特定操作
register_custom_op_symbolic('aten::asinh',  # 注意使用正確的操作名asinh_symbolic,11  # opset版本
)#register_op('asinh', asinh_symbolic, '', 9)model = Model()
input = torch.rand(1, 3, 10, 10)
torch.onnx.export(model, input, 'asinh.onnx')

這里的asinh_symbolic就是asinh的符號函數。從除g以外的第二個函數參數開始，其輸入參數應該嚴格對應它在ATen中的定義：

def asinh(input: Tensor, *, out: Optional[Tensor]=None) -> Tensor: ...

在符號函數的函數體中，g.op("Asinh", input)則完成了 ONNX 算子的定義。其中，第一個參數"Asinh"是算子在 ONNX 中的名稱。至于第二個參數?input，如我們剛剛在文檔里所見，這個算子只有一個輸入，因此我們只要把符號函數的輸入參數?input?對應過去就行。ONNX 的?Asinh?的輸出和 ATen 的?asinh?的輸出是一致的，因此我們直接把?g.op()?的結果返回即可。

定義完符號函數后，我們要把這個符號函數和原來的 ATen 算子“綁定”起來。這里，我們要用到?register_op?這個 PyTorch API 來完成綁定。如示例所示，只需要一行簡單的代碼即可把符號函數?asinh_symbolic?綁定到算子?asinh?上：?

register_op('asinh', asinh_symbolic, '', 9)

register_op第一個參數是目標ATen算子名，第二個是要注冊的符號函數，這兩個參數很好理解。第三個參數是算子的“域”，對于普通 ONNX 算子，直接填空字符串即可。第四個參數表示向哪個算子集版本注冊。我們遵照 ONNX 標準，向第 9 號算子集注冊。值得注意的是，這里向第 9 號算子集注冊，不代表較新的算子集（第 10 號、第 11 號……）都得到了注冊。在示例中，我們先只向第 9 號算子集注冊。

?測試算子

在完成了一份自定義算子后，我們一定要測試一下算子的正確性。一般我們要用 PyTorch 運行一遍原算子，再用推理引擎（比如?ONNX Runtime）運行一下 ONNX 算子，最后比對兩次的運行結果。對于我們剛剛得到的?asinh.onnx，可以用如下代碼來驗證：

import onnxruntime
import torch
import numpy as npclass Model(torch.nn.Module):def __init__(self):super().__init__()def forward(self, x):return torch.asinh(x)model = Model()
input = torch.rand(1, 3, 10, 10)
torch_output = model(input).detach().numpy()sess = onnxruntime.InferenceSession('asinh.onnx')
ort_output = sess.run(None, {'onnx::Asinh_0': input.numpy()})[0]assert np.allclose(torch_output, ort_output)

4.2 支持TorchScript算子

對于一些比較復雜的運算，僅使用 PyTorch 原生算子是無法實現的。這個時候，就要考慮自定義一個 PyTorch 算子，再把它轉換到 ONNX 中了。新增 PyTorch 算子的方法有很多，PyTorch 官方比較推薦的一種做法是添加?TorchScript 算子?。?

由于添加算子的方法較繁瑣，我們今天跳過新增 TorchScript 算子的內容，以可變形卷積（Deformable Convolution）算子為例，介紹為現有 TorchScript 算子添加 ONNX 支持的方法。

可變形卷積（Deformable Convolution）是在 Torchvision 中實現的 TorchScript 算子，雖然尚未得到廣泛支持，但是出現在許多模型中。

有了支持 ATen 算子的經驗之后，我們可以知道為算子添加符號函數一般要經過以下幾步：?

獲取原算子的前向推理接口。?
獲取目標 ONNX 算子的定義。?
編寫符號函數并綁定。?

在為可變形卷積添加符號函數時，我們也可以嘗試走一遍這個流程。

其中，torchvision.ops.DeformConv2d就是Torchvision中的可變形卷積層。相比于普通卷積，可變形卷積的其他參數都大致相同，唯一的區別就是在推理時需要多輸入一個表示偏移量的張量。

然后，我們查詢算子的前向推理接口。DeformConv2d?層最終會調用?deform_conv2d?這個算子。我們可以在?torchvision/csrc/ops/deform_conv2d.cpp?中查到該算子的調用接口：?

m.def(TORCH_SELECTIVE_SCHEMA( "torchvision::deform_conv2d(Tensor input,  Tensor weight,  Tensor offset,  ...... bool use_mask) -> Tensor"));

自定義ONNX算子?

很遺憾的是，如果我們去 ONNX 的官方算子頁面搜索 "deform"，將搜不出任何內容。目前，ONNX 還沒有提供可變形卷積的算子，我們要自己定義一個 ONNX 算子了。?

我們在前面講過，g.op()?是用來定義 ONNX 算子的函數。對于 ONNX 官方定義的算子，g.op()?的第一個參數就是該算子的名稱。而對于一個自定義算子，g.op()?的第一個參數是一個帶命名空間的算子名，比如：?

g.op("custom::deform_conv2d, ...)

其中，"::"前面的內容就是我們的命名空間。該概念和 C++ 的命名空間類似，是為了防止命名沖突而設定的。如果在?g.op()?里不加前面的命名空間，則算子會被默認成 ONNX 的官方算子。?

PyTorch 在運行?g.op()?時會對官方的算子做檢查，如果算子名有誤，或者算子的輸入類型不正確，?g.op()?就會報錯。為了讓我們隨心所欲地定義新 ONNX 算子，我們必須設定一個命名空間，給算子取個名，再定義自己的算子。?

我們在第一篇教程講過：ONNX 是一套標準，本身不包括實現。在這里，我們就簡略地定義一個 ONNX 可變形卷積算子，而不去寫它在某個推理引擎上的實現。在后續的文章中，我們再介紹在各個推理引擎中添加新 ONNX 算子支持的方法。此處，我們只關心如何導出一個包含新 ONNX 算子節點的 onnx 文件。因此，我們可以為新算子編寫如下簡單的符號函數

import torch 
import torchvision class Model(torch.nn.Module): def __init__(self): super().__init__() self.conv1 = torch.nn.Conv2d(3, 18, 3) self.conv2 = torchvision.ops.DeformConv2d(3, 3, 3) def forward(self, x): return self.conv2(x, self.conv1(x)) from torch.onnx import register_custom_op_symbolic 
from torch.onnx.symbolic_helper import parse_args @parse_args("v", "v", "v", "v", "v", "i", "i", "i", "i", "i", "i", "i", "i", "none") 
def symbolic(g,  input, weight, offset, mask, bias, stride_h, stride_w, pad_h, pad_w, dil_h, dil_w, n_weight_grps, n_offset_grps, use_mask): return g.op("custom::deform_conv2d", input, offset) register_custom_op_symbolic("torchvision::deform_conv2d", symbolic, 9) model = Model() 
input = torch.rand(1, 3, 10, 10) 
torch.onnx.export(model, input, 'dcn.onnx')

在這個符號函數中，我們以剛剛搜索到的算子輸入參數作為符號函數的輸入參數，并只用?input?和?offset?來構造一個簡單的 ONNX 算子。?

這段代碼中，最令人疑惑的就是裝飾器?@parse_args?了。簡單來說，TorchScript 算子的符號函數要求標注出每一個輸入參數的類型。比如"v"表示 Torch 庫里的?value?類型，一般用于標注張量，而"i"表示 int 類型，"f"表示 float 類型，"none"表示該參數為空。具體的類型含義可以在?torch.onnx.symbolic_helper.py?(https://github.com/pytorch/pytorch/blob/master/torch/onnx/symbolic_helper.py)中查看。這里輸入參數中的?input, weight, offset, mask, bias?都是張量，所以用"v"表示。后面的其他參數同理。我們不必糾結于?@parse_args?的原理，根據實際情況對符號函數的參數標注類型即可。

代碼成功運行的話，我們應該能得到如下的 ONNX 模型：

?4.3 使用torch.autograd.Function

最后，我們來學習一種簡單的為 PyTorch 添加 C++ 算子實現的方法，來代替較為復雜的新增 TorchScript 算子。同時，我們會用?torch.autograd.Function?封裝這個新算子。torch.autograd.Function?能完成算子實現和算子調用的隔離。不管算子是怎么實現的，它封裝后的使用體驗以及 ONNX 導出方法會和原生的 PyTorch 算子一樣。這是我們比較推薦的為算子添加 ONNX 支持的方法。?

為了應對更復雜的情況，我們來自定義一個奇怪的?my_add?算子。這個算子的輸入張量?a, b，輸出?2a + b?的值。我們會先把它在 PyTorch 中實現，再把它導出到 ONNX 中。

為PyTorch添加C++拓展

為 PyTorch 添加簡單的 C++ 拓展還是很方便的。對于我們定義的?my_add?算子，可以用以下的 C++ 源文件來實現。我們把該文件命名為 "my_add.cpp"：

// my_add.cpp #include <torch/torch.h> torch::Tensor my_add(torch::Tensor a, torch::Tensor b) 
{ return 2 * a + b; 
} PYBIND11_MODULE(my_lib, m) 
{ m.def("my_add", my_add); 
}

由于在 PyTorch 中添加 C++ 拓展和模型部署關系不大，這里我們僅給出這個簡單的示例，并不對其原理做過多講解。?

在這段代碼中，torch::Tensor?就是 C++ 中 torch 的張量類型，它的加法和乘法等運算符均已重載。因此，我們可以像對普通標量一樣對張量做加法和乘法。?

輕松地完成了算子的實現后，我們用?PYBIND11_MODULE?來為 C++ 函數提供 Python 調用接口。這里的?my_lib?是我們未來要在 Python 里導入的模塊名。雙引號中的?my_add?是 Python 調用接口的名稱，這里我們對齊 C++ 函數的名稱，依然用 "my_add"這個名字。?

之后，我們可以編寫如下的 Python 代碼并命名為 "setup.py"，來編譯剛剛的 C++ 文件：?

from setuptools import setup 
from torch.utils import cpp_extension setup(name='my_add', ext_modules=[cpp_extension.CppExtension('my_lib', ['my_add.cpp'])], cmdclass={'build_ext': cpp_extension.BuildExtension})

這段代碼使用了 Python 的 setuptools 編譯功能和 PyTorch 的 C++ 拓展工具函數，可以編譯包含了 torch 庫的 C++ 源文件。這里我們需要填寫的只有模塊名和模塊中的源文件名。我們剛剛把模塊命名為?my_lib，而源文件只有一個?my_add.cpp，因此拓展模塊那一行要寫成?ext_modules=[cpp_extension.CppExtension('my_lib', ['my_add.cpp'])],。?

之后，像處理普通的 Python 包一樣執行安裝命令，我們的 C++ 代碼就會自動編譯了。?

python setup.py develop

用 torch.autograd.Function封裝

直接用 Python 接口調用 C++ 函數不太“美觀”，一種比較優雅的做法是把這個調用接口封裝起來。這里我們用?torch.autograd.Function?來封裝算子的底層調用：

import torch 
import my_lib 
class MyAddFunction(torch.autograd.Function): @staticmethod def forward(ctx, a, b): return my_lib.my_add(a, b) @staticmethod def symbolic(g, a, b): two = g.op("Constant", value_t=torch.tensor([2])) a = g.op('Mul', a, two) return g.op('Add', a, b)

我們在前面的教程中已經見過?torch.autograd.Function，這里我們正式地對其做一個介紹。Function?類本身表示 PyTorch 的一個可導函數，只要為其定義了前向推理和反向傳播的實現，我們就可以把它當成一個普通 PyTorch 函數來使用。?

PyTorch 會自動調度該函數，合適地執行前向和反向計算。對模型部署來說，Function?類有一個很好的性質：如果它定義了?symbolic?靜態方法，該?Function?在執行?torch.onnx.export()?時就可以根據?symbolic?中定義的規則轉換成 ONNX 算子。這個?symbolic?就是前面提到的符號函數，只是它的名稱必須是?symbolic?而已。?

在?forward?函數中，我們用?my_lib.my_add(a, b)?就可以調用之前寫的C++函數了。這里?my_lib?是庫名，my_add?是函數名，這兩個名字是在前面C++的?PYBIND11_MODULE?中定義的。?

在?symbolic?函數中，我們用?g.op()?定義了三個算子：常量、乘法、加法。這里乘法和加法的用法和前面提到的?asinh?一樣，只需要根據 ONNX 算子定義規則把輸入參數填入即可。而在定義常量算子時，我們要把 PyTorch 張量的值傳入?value_t?參數中。?

在 ONNX 中，我們需要把新建常量當成一個算子來看待，盡管這個算子并不會以節點的形式出現在 ONNX 模型的可視化結果里。?

把算子封裝成?Function?后，我們可以把?my_add算子用起來了。

my_add = MyAddFunction.apply class MyAdd(torch.nn.Module): def __init__(self): super().__init__() def forward(self, a, b): return my_add(a, b)

在這份代碼里，我們先用?my_add = MyAddFunction.apply?獲取了一個奇怪的變量。這個變量是用來做什么的呢？其實，apply是torch.autograd.Function?的一個方法，這個方法完成了?Function?在前向推理或者反向傳播時的調度。我們在使用?Function?的派生類做推理時，不應該顯式地調用?forward()，而應該調用其?apply?方法。?

這里我們使用?my_add = MyAddFunction.apply?把這個調用方法取了一個更簡短的別名?my_add。以后在使用?my_add?算子時，我們應該忽略?MyAddFunction?的實現細節，而只通過?my_add?這個接口來訪問算子。這里?my_add?的地位，和 PyTorch 的?asinh,?interpolate,?conv2d等原生函數是類似的。?

有了訪問新算子的接口后，我們可以進一步把算子封裝成一個神經網絡中的計算層。我們定義一個叫做的?MyAdd?的?torch.nn.Module，它封裝了my_add，就和封裝了conv2d?的?torch.nn.Conv2d?一樣。

測試算子

費了好大的功夫來“包裝”我們的新算子后，我們終于可以來使用它了。和之前的測試流程一樣，讓我們用下面的代碼來導出一個包含新算子的 ONNX 模型，并驗證一下它是否正確。?

model = MyAdd() 
input = torch.rand(1, 3, 10, 10) 
torch.onnx.export(model, (input, input), 'my_add.onnx') 
torch_output = model(input, input).detach().numpy() import onnxruntime 
import numpy as np 
sess = onnxruntime.InferenceSession('my_add.onnx') 
ort_output = sess.run(None, {'a': input.numpy(), 'b': input.numpy()})[0] assert np.allclose(torch_output, ort_output)

在這份代碼中，我們直接把?MyAdd?作為要導出的模型。我們計算了一個 PyTorch 模型的運行結果，又導出 ONNX 模型，計算了 ONNX 模型在 ONNX Runtime 上的運算結果。如果一切正常的話，這兩個結果是一樣的，這份代碼不會報任何錯誤，沒有任何輸出。

可視化一下?my_add.onnx，可以看出，和我們設計得一樣，my_add?算子被翻譯成了兩個 ONNX 算子節點（其中常量算子被放入了?Mul?的參數中）。?

整理一下，整個流程的 Python 代碼如下：

import torch 
import my_lib 
class MyAddFunction(torch.autograd.Function): @staticmethod def forward(ctx, a, b): return my_lib.my_add(a, b) @staticmethod def symbolic(g, a, b): two = g.op("Constant", value_t=torch.tensor([2])) a = g.op('Mul', a, two) return g.op('Add', a, b) my_add = MyAddFunction.apply class MyAdd(torch.nn.Module): def __init__(self): super().__init__() def forward(self, a, b): return my_add(a, b) model = MyAdd() 
input = torch.rand(1, 3, 10, 10) 
torch.onnx.export(model, (input, input), 'my_add.onnx') 
torch_output = model(input, input).detach().numpy() import onnxruntime 
import numpy as np 
sess = onnxruntime.InferenceSession('my_add.onnx') 
ort_output = sess.run(None, {'a': input.numpy(), 'b': input.numpy()})[0] assert np.allclose(torch_output, ort_output)

總結?

ATen 是 PyTorch 的 C++ 張量運算庫。通過查詢?torch/_C/_VariableFunctions.pyi?和?torch/nn/functional.pyi，我們可以知道 ATen 算子的 Python 接口定義。?
用?register_custom_op_symbolic?可以為 ATen 算子補充注冊符號函數?
用?register_custom_op_symbolic?可以為 TorchScript 算子補充注冊符號函數?
如何在 PyTorch 里添加 C++ 拓展?
如何用?torch.autograd.Function?封裝一個自定義 PyTorch 算子?
如何編寫符號函數?symbolic(g, ...)。?
如何用?g.op()?把一個 PyTorch 算子映射成一個或多個 ONNX 算子，或者是自定義的 ONNX 算子。

五. ONNX 模型的修改與調試

不知道大家會不會有這樣一些疑問：ONNX 模型在底層是用什么格式存儲的？如何不依賴深度學習框架，只用 ONNX 的 API 來構造一個 ONNX 模型？如果沒有源代碼，只有一個 ONNX 模型，該如何對這個模型進行調試？別急，今天我們就來為大家一一揭曉。?
在這期教程里，我們將圍繞 ONNX 這一套神經網絡定義標準本身，探究 ONNX 模型的構造、讀取、子模型提取、調試。首先，我們會學習 ONNX 的底層表示方式。之后，我們會用 ONNX API 構造和讀取模型。最后，我們會利用 ONNX 提供的子模型提取功能，學習如何調試 ONNX 模型。

5.1 ?ONNX的底層實現

1. ONNX的存儲格式

ONNX在底層時用Protobuf定義的。Protobuf,全稱Protocol Buffer，是Google提出的一套表示和序列化數據的機制。使用protobuf時，用戶需要先寫一份數據定義文件，再根據這份定義文件把數據存儲進一份二進制文件。可以說，數據定義文件就是數據類，二進制文件就是數據類的實例。

這里給出一個 Protobuf 數據定義文件的例子：?

message Person { required string name = 1; required int32 id = 2; optional string email = 3; 
}

這段定義表示在Person這種數據類型中，必須包含 name、id這兩個字段，選擇性包含email字段。根據這份定義文件，用戶就可以選擇一種編程語言，定義一個含有成員變量name、id、email的Person類，把這個類的某個實例用Protobuf存儲成二進制文件；反之，用戶也可以用二進制文件和對應的數據定義文件，讀取出一個Person類的實例。而對于ONNX，Protobuf的數據定義文件在其開源庫，這些文件定義了神經網絡中模型、節點、張量的數據類型規范；而二進制文件就是我們熟悉的“.onnx”文件，每一個onnx文件按照數據定義規范，存儲了一個神經網絡的所有相關數據。直接用protobuf生成ONNX模型還是比較麻煩。幸虧的事，ONNX提供了很多使用API，我們可以在完全不了解Protobuf前提下，構造和讀取ONNX模型。

2. ONNX的結構定義

再用API對ONNX模型操作之前，我們還需要先了解一下ONNX的結構定義規則，學習一下ONNX在Protobuf定義文件里事怎樣描述一個神經網絡的。

回想一下，神經網絡本質上是一個計算圖，計算圖的節點是算子，邊是參與運算的張量。而通過可視化ONNX模型，我們知道ONNX記錄了所有算子節點的屬性信息，并把參與運算的張量信息存儲在算子節點的輸入輸出信息中。實際上，ONNX模型的結構可以用類圖大致表示如下：

如圖所示，一個 ONNX 模型可以用?ModelProto?類表示。ModelProto?包含了版本、創建者等日志信息，還包含了存儲計算圖結構的?graph。GraphProto?類則由輸入張量信息、輸出張量信息、節點信息組成。張量信息?ValueInfoProto?類包括張量名、基本數據類型、形狀。節點信息?NodeProto?類包含了算子名、算子輸入張量名、算子輸出張量名。?
讓我們來看一個具體的例子。假如我們有一個描述?output=a*x+b?的 ONNX 模型?model，用?print(model)?可以輸出以下內容：

ir_version: 8 
graph { node { input: "a" input: "x" output: "c" op_type: "Mul" } node { input: "c" input: "b" output: "output" op_type: "Add" } name: "linear_func" input { name: "a" type { tensor_type { elem_type: 1 shape { dim {dim_value: 10} dim {dim_value: 10} } } } } input { name: "x" type { tensor_type { elem_type: 1 shape { dim {dim_value: 10} dim {dim_value: 10} } } } } input { name: "b" type { tensor_type { elem_type: 1 shape { dim {dim_value: 10} dim {dim_value: 10} } } } } output { name: "output" type { tensor_type { elem_type: 1 shape { dim { dim_value: 10} dim { dim_value: 10} } } } } 
} 
opset_import {version: 15}

對應上文中的類圖，這個模型的信息由ir_version, opset_import等全局信息和graph圖信息組成。而graph包含一個乘法節點、一個加法節點、三個輸入張量a,x,b以及一個輸出張量output。在下一節里，我們會會API構造出這個模型，并輸出這段結果。

3. 讀寫ONNX模型

3.1 構造ONNX模型

在上一小節中，我們知道了 ONNX 模型是按以下的結構組織起來的：?

ModelProto?
- GraphProto?
  - NodeProto?
  - ValueInfoProto

現在，讓我們拋開PyTorch,嘗試完全用ONNX的PyThon API構造一個描述線性函數output=a*x +b的ONNX模型。我們將根據上面的結構，自底向上地構造這個模型。首先我們可以用helper.make_tensor_value_info構造出一個描述張量信息的ValueInfoProto對象。如前面的類圖所示，我們要傳入張量，他們的表示方法都是一樣的。因此，這里我們用類似的方式為三個輸入a,x,b和一個輸出output構造ValueInfoProto對象。如下：

import onnx
from onnx import helper
from onnx import TensorProtoa = helper.make_tensor_value_info('a', TensorProto.FLOAT, [10, 10])
x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10, 10])
b = helper.make_tensor_value_info('b', TensorProto.FLOAT, [10, 10])
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [10, 10])

之后，我們要構造算子節點信息NodeProto,這可以通過在helper.make_node中傳入算子類型、輸入算子名、輸出算子名這三個信息來實現。我們這里先構造了描述c=a*x的乘法節點，再構造了output=c+b的加法節點。如下：

mul = helper.make_node('Mul', ['a', 'b'], ['c'])
add = helper.make_node('Add', ['c', 'b'], ['output'])

在計算機中，圖一般是用一個節點集和一個邊集表示的。而ONNX巧妙地把邊的信息保存在了節點信息里，省去了保存邊集的步驟。在ONNX，如果某節點的輸入名和之前某節點的輸出名相同，就默認了這兩個節點是相連的。如上圖例子所示：Mul節點定義了輸出c, Add節點定義了輸入c, 則Mul節點和Add節點是相連的。

正是因為這種邊的隱式定義規則，所以ONNX對節點的輸入由一定的要求：一個節點的輸入，要么是整個模型的輸入，要么是之前某個節點的輸出。如果我們把a, x, b中的某個輸入節點從計算圖中拿出，或者把Mul的輸出從c改成d,則最終的ONNX模型都是不滿足標準的。

一個不滿足標準的 ONNX 模型可能無法被推理引擎正確識別。ONNX 提供了 API?onnx.checker.check_model?來判斷一個 ONNX 模型是否滿足標準。

接下來，我們用helper.make_graph來構造計算圖GraphProto。helper.make_graph函數需要傳入節點、圖名稱、輸入張量信息、輸出張量信息這4個參數。如下面的代碼所示，我們把之前構造出來的NodeProto對象和ValueInfoProto對象按照順序傳入即可。?

graph = helper.make_graph([mul, add], 'linear_func', [a, x, b], [output])

這里make_graph的節點參數有一個要求：計算圖的節點必須以拓撲序給出。

我們以剛剛構造出來的這個計算圖為研究對象，通過下圖展示的兩個例子來直觀理解拓撲序。

這里我們只關注?Mul?和?Add?節點以及它們之間的邊?c。在情況 1 中：如果我們的節點以?[Mul, Add]?順序給出，那么遍歷到?Add?時，它的輸入?c?可以在之前的Mul的輸出中找到。但是，如情況 2 所示：如果我們的節點以?[Add, Mul]?的順序給出，那么?Add?就找不到輸入邊，計算圖也無法成功構造出來了。這里的?[Mul, Add]?就是符合有向圖的拓撲序的，而?[Add, Mul]?則不滿足。

?最后，我們用helper.make_model把計算圖GraphProto封裝進模型ModelProto里，一個ONNX模型就構造完成了。make_model函數中還可以添加模型制作者、版本等信息。

model = helper.make_model(graph)

構造完模型之后，我們用下面這三行代碼來檢查模型正確性、把模型以文本形式輸出、存儲到一個 ".onnx" 文件里。這里用?onnx.checker.check_model?來檢查模型是否滿足 ONNX 標準是必要的，因為無論模型是否滿足標準，ONNX 都允許我們用?onnx.save?存儲模型。我們肯定不希望生成一個不滿足標準的模型。?

onnx.checker.check_model(model) 
print(model) 
onnx.save(model, 'linear_func.onnx')

成功執行這些代碼的話，程序會以文本格式輸出模型的信息，其內容應該和我們在上一節展示的輸出一樣。?
整理一下，用 ONNX Python API 構造模型的代碼如下：

import onnx
from onnx import helper
from onnx import TensorProtoa = helper.make_tensor_value_info('a', TensorProto.FLOAT, [10, 10])
x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10, 10])
b = helper.make_tensor_value_info('b', TensorProto.FLOAT, [10, 10])
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [10, 10])mul = helper.make_node('Mul', ['a', 'x'], ['c'])
add = helper.make_node('Add', ['c', 'b'], ['output'])graph = helper.make_graph([mul, add], 'linear_func', [a, x, b], [output])
model = helper.make_model(graph)
onnx.checker.check_model(model) 
print(model) 
onnx.save(model, 'linear_func.onnx')

?老規矩，我們可以用?ONNX Runtime?運行模型，來看看模型是否正確：

import onnxruntime 
import numpy as np sess = onnxruntime.InferenceSession('linear_func.onnx') 
a = np.random.rand(10, 10).astype(np.float32) 
b = np.random.rand(10, 10).astype(np.float32) 
x = np.random.rand(10, 10).astype(np.float32) output = sess.run(['output'], {'a': a, 'b': b, 'x': x})[0] assert np.allclose(output, a * x + b)

一切順利的話，這段代碼不會有任何報錯信息。這說明我們的模型等價于執行?a * x + b?這個計算。

3.2 讀取并修改ONNX模型

通過用 API 構造 ONNX 模型，我們已經徹底搞懂了 ONNX 由哪些模塊組成。現在，讓我們看看該如何讀取現有的".onnx"文件并從中提取模型信息。?
首先，我們可以用下面的代碼讀取一個 ONNX 模型：

import onnx 
model = onnx.load('linear_func.onnx') 
print(model)

之前在輸出模型時，我們傳給onnx.save的是一個ModelProto對象。同理，我們用上面的onnx.load讀取ONNX模型時，我們收獲的也是一個Model.Proto對象。輸出這個對象后，我們應該得到和之前完全相同的輸出。

接下來，我們來看看怎么把圖GraphProto、節點?NodeProto、張量信息?ValueInfoProto?讀取出來：

graph = model.graph 
node = graph.node 
input = graph.input 
output = graph.output 
print(node) 
print(input) 
print(output)

使用如上這些代碼，我們可以分別訪問模型的圖、節點、張量信息。這里大家或許會有疑問：該怎樣找出?graph.node,graph.input?中?node, input?這些屬性名稱呢？其實，屬性的名稱就寫在每個對象的輸出里。我們以?print(node)?的輸出為例：

[input: "a"

input: "x"

output: "c"

op_type: "Mul" ,

input: "c"

input: "b"

output: "output"

op_type: "Add" ]

在這段輸出中，我們能看出?node?其實就是一個列表，列表中的對象有屬性?input, output, op_type（這里?input?也是一個列表，它包含的兩個元素都顯示出來了）。我們可以用下面的代碼來獲取?node?里第一個節點?Mul?的屬性：

node_0 = node[0]
node_0_inputs = node_0.inputnode_0_outputs = node_0.outputinput_0 = node_0_inputs[0]
print(input_0) #a
input_1 = node_0_inputs[1]
print(input_1) #x
output = node_0_outputs[0]
print(output) #c
op_type = node_0.op_type
print(op_type) #Mul

當我們想知道 ONNX 模型某數據對象有哪些屬性時，我們不必去翻 ONNX 文檔，只需要先把數據對象輸出一下，然后在輸出結果找出屬性名即可。?
讀取 ONNX 模型的信息后，修改 ONNX 模型就是一件很輕松的事了。我們既可以按照上一小節的模型構造方法，新建節點和張量信息，與原有模型組合成一個新的模型，也可以在不違反 ONNX 規范的前提下直接修改某個數據對象的屬性。?
這里我們來看一個直接修改模型屬性的例子：

import onnx 
model = onnx.load('linear_func.onnx') node = model.graph.node 
node[1].op_type = 'Sub' onnx.checker.check_model(model) 
onnx.save(model, 'linear_func_2.onnx')

在讀入之前的?linear_func.onnx?模型后，我們可以直接修改第二個節點的類型?node[1].op_type，把加法變成減法。這樣，我們的模型描述的是?a * x - b?這個線性函數。大家感興趣的話，可以用 ONNX Runtime 運行新模型?linear_func_2.onnx，來驗證一下它和?a * x - b?是否等價。

4. 調試ONNX模型

在實際部署中，如果用深度學習框架導出的 ONNX 模型出了問題，一般要通過修改框架的代碼來解決，而不會從 ONNX 入手，我們把 ONNX 模型當成一個不可修改的黑盒看待。?
現在，我們已經深入學習了 ONNX 的原理，可以嘗試對 ONNX 模型本身進行調試了。在這一節里，讓我們看看該如何巧妙利用 ONNX 提供的子模型提取功能，對 ONNX 模型進行調試。

4.1 子模型提取

ONNX 官方為開發者提供了子模型提取（extract）的功能。子模型提取，顧名思義，就是從一個給定的 ONNX 模型中，拿出一個子模型。這個子模型的節點集、邊集都是原模型中對應集合的子集。讓我們來用 PyTorch 導出一個復雜一點的 ONNX 模型，并在它的基礎上執行提取操作：

import torch class Model(torch.nn.Module): def __init__(self): super().__init__() self.convs1 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3), torch.nn.Conv2d(3, 3, 3), torch.nn.Conv2d(3, 3, 3)) self.convs2 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3), torch.nn.Conv2d(3, 3, 3)) self.convs3 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3), torch.nn.Conv2d(3, 3, 3)) self.convs4 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3), torch.nn.Conv2d(3, 3, 3), torch.nn.Conv2d(3, 3, 3)) def forward(self, x): x = self.convs1(x) x1 = self.convs2(x) x2 = self.convs3(x) x = x1 + x2 x = self.convs4(x) return x model = Model() 
input = torch.randn(1, 3, 20, 20) torch.onnx.export(model, input, 'whole_model.onnx')

在前面的章節中，我們學過，ONNX 的邊用同名張量表示的。也就是說，這里的邊序號，實際上是前一個節點的輸出張量序號和后一個節點的輸入張量序號。由于這個模型是用 PyTorch 導出的，這些張量序號都是 PyTorch 自動生成的。

接著，我們可以下面的代碼提取出一個子模型：?

import onnx  onnx.utils.extract_model('whole_model.onnx', 'partial_model.onnx', ['/convs1/convs1.2/Conv_output_0'], ['/convs4/convs4.0/Conv_output_0'])

?paritial mode的子模型可視化結果如下圖：

通過觀察代碼和輸出圖，應該不難猜出這段代碼的作用是把原計算圖從邊 22 到邊 28 的子圖提取出來，并組成一個子模型。onnx.utils.extract_model?就是完成子模型提取的函數，它的參數分別是原模型路徑、輸出模型路徑、子模型的輸入邊（輸入張量）、子模型的輸出邊（輸出張量）。?
直觀地來看，子模型提取就是把輸入邊到輸出邊之間的全部節點都取出來。那么，這個功能在使用上有什么限制呢？基于?whole_model.onnx, 我們來看一看三個子模型提取的示例。

添加額外輸出

我們在提取時新設定了一個輸出張量，如下面的代碼所示：

onnx.utils.extract_model('whole_model.onnx', 'submodel_1.onnx', ['/convs1/convs1.2/Conv_output_0'], ['/convs2/convs2.1/Conv_output_0', '31'])

添加冗余輸入?

如果我們還是像開始一樣提取邊 22 到邊 28 之間的子模型，但是多添加了一個輸入?input.1，那么提取出的子模型會有一個冗余的輸入?input.1，如下面的代碼所示：?

onnx.utils.extract_model('whole_model.onnx', 'submodel_2.onnx', ['22', 'input.1'], ['28'])

從下圖可以看到：無論給這個輸入傳入什么值，都不會影響子模型的輸出。可以認為如果只用子模型的部分輸入就能得到輸出，那么那些”較早“的多出來的輸入就是冗余的。

輸入信息不足

這次，我們嘗試提取的子模型輸入是邊 24，輸出是邊 28。如下面的代碼和圖所示：

#error
onnx.utils.extract_model('whole_model.onnx', 'submodel_3.onnx', ['24'], ['28'])

從圖中可以看出，想通過邊 24 計算邊 28 的結果，至少還需要輸入邊 26，或者更上面的邊。僅憑借邊 24 是無法計算出邊 28 的結果的，因此這樣提取子模型會報錯。?

通過上面幾個使用示例，我們可以整理出子模型提取的實現原理：新建一個模型，把給定的輸入和輸出填入。之后把圖的所有有向邊反向，從輸出邊開始遍歷節點，碰到輸入邊則停止，把這樣遍歷得到的節點做為子模型的節點。?
如果還沒有徹底弄懂這個提取原理，沒關系，我們只要盡量保證在填寫子模型的輸入輸出時，讓輸出恰好可以由輸入決定即可。?

5. 輸出 ONNX 中間節點的值?

在使用 ONNX 模型時，最常見的一個需求是能夠用推理引擎輸出中間節點的值。這多見于深度學習框架模型和 ONNX 模型的精度對齊中，因為只要能夠輸出中間節點的值，就能定位到精度出現偏差的算子。我們來看看如何用子模型提取實現這一任務。?
在剛剛的第一個子模型提取示例中，我們添加了一條原來模型中不存在的輸出邊。用同樣的原理，我們可以在保持原有輸入輸出不變的同時，新增加一些輸出，提取出一個能輸出中間節點的”子模型“。例如：

onnx.utils.extract_model('whole_model.onnx', 'more_output_model.onnx', ['input.1'], ['31','/convs3/convs3.1/Conv_output_0', '/convs2/convs2.1/Conv_output_0'])

?這樣，用 ONNX Runtime 運行?more_output_model.onnx?這個模型時，我們就能得到更多的輸出了。?
為了方便調試，我們還可以把原模型拆分成多個互不相交的子模型。這樣，在每次調試時，可以只對原模型的部分子模塊調試。比如：


onnx.utils.extract_model('whole_model.onnx', 'debug_model_1.onnx', ['input.1'], ['/convs2/convs2.0/Conv_output_0'])
onnx.utils.extract_model('whole_model.onnx', 'debug_model_2.onnx', ['/convs1/convs1.1/Conv_output_0'], ['/convs3/convs3.1/Conv_output_0']) 
onnx.utils.extract_model('whole_model.onnx', 'debug_model_3.onnx', ['/convs1/convs1.2/Conv_output_0'], ['/convs2/convs2.1/Conv_output_0']) 
onnx.utils.extract_model('whole_model.onnx', 'debug_model_4.onnx', ['/convs3/convs3.1/Conv_output_0', '/convs2/convs2.1/Conv_output_0'],['31'])

子模型提取固然是一個便利的 ONNX 調試工具。但是，在實際的情況中，我們一般是用 PyTorch 等框架導出 ONNX 模型。這里有兩個問題：?

一旦 PyTorch 模型改變，ONNX 模型的邊序號也會改變。這樣每次提取同樣的子模塊時都要重新去 ONNX 模型里查序號，如此繁瑣的調試方法是不會在實踐中采用的。?
即使我們能保證 ONNX 的邊序號不發生改變，我們也難以把 PyTorch 代碼和 ONNX 節點對應起來——當模型結構變得十分復雜時，要識別 ONNX 中每個節點的含義是不可能的。?

在 MMDeploy 中，我們為 PyTorch 模型添加了模型分塊功能。使用這個功能，我們可以通過只修改 PyTorch 模型的實現代碼來把原模型導出成多個互不相交的子 ONNX 模型。我們會在后續教程中對其介紹。

總結?

在這篇教程中，我們拋開了 PyTorch，學習了 ONNX 模型本身的知識。老規矩，我們來總結一下這篇教程的知識點：?

ONNX 使用 Protobuf 定義規范和序列化模型。?
一個 ONNX 模型主要由?ModelProto,GraphProto,NodeProto,ValueInfoProto?這幾個數據類的對象組成。?
使用?onnx.helper.make_xxx，我們可以構造 ONNX 模型的數據對象。?
onnx.save()?可以保存模型，onnx.load()?可以讀取模型，onnx.checker.check_model()?可以檢查模型是否符合規范。?
onnx.utils.extract_model()?可以從原模型中取出部分節點，和新定義的輸入、輸出邊構成一個新的子模型。?
利用子模型提取功能，我們可以輸出原 ONNX 模型的中間結果，實現對 ONNX 模型的調試。?

至此，我們對 ONNX 相關知識的學習就告一段落了。回顧一下，我們先學習了 PyTorch 轉 ONNX 有關 API 的用法；接著，我們學習了如何用自定義算子解決 PyTorch 和 ONNX 表達能力不足的問題；最后我們單獨學習了 ONNX 模型的調試方法。通過對 ONNX 由淺入深的學習，我們基本可以應對模型部署中和 ONNX 有關的絕大多數問題了。?

六.?實現 PyTorch-ONNX 精度對齊工具

精度對齊，是模型部署中重要的一個環節。把深度學習框架模型轉換成中間表示模型后，部署工程師們要做的第一件事就是精度對齊，確保模型的計算結果與之前相當。精度對齊時最常用的方法就是使用測試集評估一遍中間表示模型，看看模型的評估指標準確度和相似度是否下降。

而在 PyTorch 到 ONNX 這條部署路線上，這種精度對齊方式有一些不便：一旦我們發現 PyTorch 模型和 ONNX 模型的評估指標有了出入，我們很難去追蹤精度是在哪一個模塊出了問題。這是因為 PyTorch 和 ONNX 模塊總是難以對應。如下面的例子所示：

假設我們現在有一個由很多卷積塊?convs1, convs2...?組成的網絡，我們想對齊 PyTorch 模型和 ONNX 模型的精度。第一步，我們想比較第一個卷積塊的輸出?x = self.convs1(x)。模塊在PyTorch 模型中的輸出可以很輕松地得到，可是，這個輸出究竟對應 ONNX 模型里的哪一個輸出呢？在小模型里，我們或許能夠通過閱讀 PyTorch 模型的源碼，推斷出每個 ONNX 模塊與 PyTorch 模塊的對應關系；但是，在大模型中，我們是難以建立 PyTorch 與 ONNX 的對應關系的。?

在這篇教程中，我們就來利用之前學過的自定義算子、子模型提取等工具，實現一個簡單的 PyTorch-ONNX 精度對齊工具。

6.1 設計思路

為了把 PyTorch 和 ONNX 模塊對應起來，我們可以使用一種儲存了調試信息的自定義算子，如下圖所示：

我們可以定義一個叫做Debug的ONNX算子，它有一個屬性調試名name。而由于每一個ONNX算子節點又自帶了輸出張量的名稱，這樣一來，ONNX節點的輸出名和調試名綁定在了一起。我們可以順著PyTorch里調試名，找到對應ONNX里的輸出，完成PyTorch和ONNX的對應。

比如在上圖的例子中，我們把第一個卷積塊輸出x=self.convs1(x)接入一個帶有調試名x_0的調試算子。在最后生成的ONNX模型中，假設調試名x_0對應的輸出張量叫做a。知道了這一信息后，我們只需要先運行一遍 PyTorch 模型，記錄第一個卷積塊的輸出；再運行一遍 ONNX 模型，用上篇教程中提到的截取 ONNX 中間結果的方法，記錄中間張量?a?的值。這樣，我們就可以對齊某 PyTorch 模塊和它對應的 ONNX 模塊的輸出了。

6.2 代碼實現

debug算子

首先，我們需要實現之前提到的Debug算子：

import torch
class DebugOp(torch.autograd.Function):@staticmethoddef forward(ctx, x, name):return x@staticmethoddef symbolic(g, x, name):return g.op("my::Debug", x, name_s=name)debug_apply = DebugOp.apply

?Debug 算子的調用接口有兩個參數：輸入張量?x?和調試名?name。為了把這個算子“偽裝”成一個普通的算子，使之能正常地參與推理、構建計算圖的操作，我們還是需要正確定義對輸入?x進行操作的?forward?函數。而在表示 PyTorch 與 ONNX 映射規則的?symbolic?函數里，我們要定義一個帶有調試名的 ONNX 算子，并把輸入的?name?傳給算子。

由于 Debug 算子本身不表示任何計算，因此在?forward?函數中，直接把輸入?x?返回即可。?

而?symbolic?函數定義了一個新算子?my::Debug：算子有一個輸入?x，一個屬性?name。我們直接把算子調用接口里的?x，name?傳入即可。

這里需要補充介紹算子定義函數?g.op()?的一些規范。在g.op()中，算子的屬性需要以?{attibute_name}_{type}=attibute_value?這樣的格式傳入。其中?{attibute_name}為屬性名，{type}?指定了算子屬性的數據類型。比如說我們上面的算子屬性寫成?name_s，實際上是定義了一個字符串類型，名字叫做?name?的屬性。除了表示字符串類型的?_s?外，還有表示?float?型的?_f，表示?tensor?型的?_t。?

在完成算子的定義后，我們可以通過?debug_apply = DebugOp.apply?獲取算子的調用接口。這樣以后就可以通過?debug_apply(x, name)?來使用這個算子了。

Debugger類

接著，我們來實現精度對齊工具的核心——Debugger 類。這個類包含了實現精度對齊所需的所有操作。其定義如下：

?Debugger類有三個成員變量：

torch_value 記錄了運行PyTorch模型后每個調試張量的值
onnx_value 記錄了運行ONNX模型后每個調試張量的值
output_debug_name：記錄了把調試張量加入ONNX的輸出后，每個輸出張量的調試名

稍后我們會在類實現的代碼中看到這些成員變量的具體用法。

Debugger類有以下方法：

debug封裝了之前變好的debug_apply。該方法需要在原PyTorch模型中調用，可以為導出的ONNX的模型添加Debug算子節點，同時記錄PyTorch調試張量值。
extract_debug_model和ONNX的子模型提取函數的用法類似，可以把帶調試節點的ONNX模型轉化成一個可以輸出調試張量的ONNX模型。
run_debug_model會使用ONNX Runtime運行模型，得到ONNX調試張量值。
print_debug_result會比較PyTorch和ONNX的調試張量值，輸出比較的結果

這4個方法會一次被調用：

生成調試節點

def debug(self, x, name):self.torch_value[name] = x.detach().cpu().numpy()return debug_apply(x, name)

如前文所述，debug完成兩件事：記錄PyTorch模型中的調試張量的值、添加Debug節點。我們使用self.torch_value[name]=x.detach().cpu().numpy()把調試張量轉成numpy格式并保存進torch_value詞典里。之后，我們調用之前編寫的debug_apply算子。

提取調試模型

import onnx
def extract_debug_model(self, input_path, output_path):model = onnx.load(input_path)inputs = [input.name for input in model.graph.input]outputs = []for node in model.graph.node:if node.op_type == 'Debug':#記錄調試張量名debug_name = node.attribute[0].s.decode('ASCII')self.output_debug_name.append(debug_name)#添加輸入output_name = node.output[0]outputs.append(output_name)#轉換Debug 節點為Indentity節點node.np_type = 'Identity'node.domain = ''del node.attribute[:]e = onnx.util.Extractor(model)extracted = e.extrac_model(inputs, outputs)onnx.save(extracted, output_path)

在PyTorch模型中插入debug方法后，我們可以得到一個包含了若干Debug節點的ONNX模型。但是這個ONNX模型不是我們最終拿來執行的模型。為了得到Debug節點的輸出（即調試張量的值），我們需要做三項處理以提取出一個可運行的調試模型：

記錄每個調試張量的調試名，為之后對齊PyTorch、ONNX調試張量值做準備。
把所有Debug節點的輸出加入到整個模型的輸出中，這樣在運行模型后就能得到這些中間節點的輸出了。
自定義的Debug節點在推理引擎中時沒有實現的，為了讓處理湖的ONNX模型運行起來，需要把Debug節點轉化成可運行的Identity（恒等）節點。

完成了這三項處理后，我們才能進行模型提取。下面，我們來看看模型提取和這幾項處理是怎么實現的。

首先，看一下和模型提取有關的代碼：?

model = onnx.load(input_path) 
inputs = [input.name for input in model.graph.input] 
outputs = [] # 獲取 outputs 
... # 調用提取模型 API 
e = onnx.utils.Extractor(model) 
extracted = e.extract_model(inputs, outputs) # 保存模型 
onnx.save(extracted, output_path)

在提取模型時，我們要準備新模型的輸入和輸出。輸入張量inputs還是保持原狀，而輸出張量outputs會在之后填入Debug節點的輸出。獲取完outputs后，我們調用提取模型的API，得到處理過后的模型，并保存此模型。

接著，看一下主處理邏輯：?

for node in model.graph.node: if node.op_type == 'Debug': ...

為了獲取和Debug節點相關的信息，我們需要遍歷ONNX模型的所有節點，找出那些類型為Debug的節點，對這些節點執行操作。

下面的代碼實現了記錄調試張量名：?

debug_name = node.attribute[0].s.decode('ASCII') 
self.output_debug_name.append(debug_name)

這段代碼的作用是：從節點的第一個屬性（即name）中取出調試名信息，并存入output_debug_name中。節點第一個屬性的值可以通過node.attribute[0]獲得。由于name是屬性是字符串，這里要用.s獲取屬性的字符串值。又由于ONNX是以二進制的形式保存所有數據的，這里要用.decode(‘ASCII’）把二進制字符串轉成一個文本字符串。

接下來的代碼用于填寫新模型輸出outputs:

output_name = node.output[0] 
outputs.append(output_name)

node.output[0]就是debug節點的輸出張量在ONNX里的名稱，把這個名稱加入新模型的輸出后，只需要運行新模型，就可以得到該輸出張量的值了。

最后這段代碼用于更改Debug節點的類型：

node.op_type = 'Identity' 
node.domain = '' 
del node.attribute[:]

為了消除 ONNX 不支持的 Debug 節點，一種比較簡單的方式是直接把 Debug 節點修改成不執行任何操作的?Indentity?類型的節點。為了做這個轉換，我們要先修改節點類型名?node.op_type?為Identity，再把節點的域（即命名空間）node.domain?修改成空，最后刪除節點的所有屬性，保證節點符合 ONNX 的規范。?
回憶一下，如果一個節點的?domain?為空，這個節點就會被當成一個 ONNX 原生算子節點。

運行調試模型

在生成調試節點時，我們已經順便記錄了Pytorch模型調試張量的值，下一步，我們要運行調試模型，記錄ONNX模型調試張量的值。實現如下：

import onnxruntime
def run_debug_model(self, input, debug_model):sess = onnxruntime.InferenceSession(debug_model, providers=['CPUExecutionProvider'])onnx_outputs = sess.run(None, input)for name, value in zip(self.output_debug_name, onnx_outputs):self.onnx_value[name] = value

?在運行調試模型前，我們要給出模型輸入、模型名這兩個參數。根據這些參數，run_debug_model會調用ONNX runtime的API，對ONNX模型進行推理。在得到了ONNX模型的輸出后，用使用上一部得到的output_debug_name信息，填寫onnx_value，把ONNX中間運算結果綁定到調試名上。完成這些步驟之后，我們就有足夠的信息做精度對齊了。

def print_debug_result(self):for name in self.torch_value.keys():if name in self.onnx_value:mse = np.mean((self.torch_value[name] - self.onnx_value[name])**2)

最后，我們同時遍歷?self.torch_value?和?self.onnx_value?這兩個詞典，比較同一個張量在 PyTorch 模型和 ONNX 模型里的輸出。在循環體中，我們只需要使用?self.torch_value[name]?和?self.onnx_value[name]?就可以訪問同一個張量在 PyTorch 里的值和在 ONNX 里的值。作為示例，這里我們可以計算二者的均方誤差?mse，以此為精度對齊的依據。

使用方法

實現了精度對齊工具后，我們來看看該怎么把這個工具用起來。?

現在，假設我們得到了一個這樣的模型：?

class Model(torch.nn.Module): def __init__(self): super().__init__() self.convs1 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3, 1, 1), torch.nn.Conv2d(3, 3, 3, 1, 1), torch.nn.Conv2d(3, 3, 3, 1, 1)) self.convs2 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3, 1, 1), torch.nn.Conv2d(3, 3, 3, 1, 1)) self.convs3 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3, 1, 1), torch.nn.Conv2d(3, 3, 3, 1, 1)) self.convs4 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3, 1, 1), torch.nn.Conv2d(3, 3, 3, 1, 1), torch.nn.Conv2d(3, 3, 3, 1, 1)) def forward(self, x): x = self.convs1(x) x = self.convs2(x) x = self.convs3(x) x = self.convs4(x) return x torch_model = Model()

沒錯！這就是本文開頭展示的那個全卷積網絡。現在我們想對齊?convs1?至?convs4?這每一個卷積塊的輸出精度，該怎么使用之前寫好的精度對齊工具呢？?

首先，我們生成管理類?Debugger?的一個實例：?

debugger = Debugger()

之后，我們要設法把 Debug 節點插入原模型：?

from types import MethodType def new_forward(self, x): x = self.convs1(x) x = debugger.debug(x, 'x_0') x = self.convs2(x) x = debugger.debug(x, 'x_1') x = self.convs3(x) x = debugger.debug(x, 'x_2') x = self.convs4(x) x = debugger.debug(x, 'x_3') return x torch_model.forward = MethodType(new_forward, torch_model)

我們可以為原模型新寫一個?forward?函數。在這個新的函數函數中，我們可以通過?debugger.debug?把每一個輸出張量標記起來，并各取一個不重復的調試名。

有了?new_forward?函數，我們需要使用?MethodType?這個 Python API 把這個函數變成模型實例?torch_model?的一個成員方法，確保?torch_model?的?forward?函數能夠被正確替換。

實現了”貍貓換太子“般巧妙的操作后，我們就可以使用 PyTorch API 導出一個帶有 Debug 節點的 ONNX 模型了：?

dummy_input = torch.randn(1, 3, 10, 10) 
torch.onnx.export(torch_model, dummy_input, 'before_debug.onnx', input_names=['input'])

由于?torch.onnx.export?模型使用的是跟蹤法，模型的?forward?函數會被執行一次，?debugger.debug?操作可以把 PyTorch 模型的調試張量輸出記錄在?debugger.torch_value?里。?

這個?before_debug.onnx?模型的部分可視化結果如下:

接下來，我們替換掉所有 Debug 節點，并記錄每個 Debug 輸出張量的 ONNX 名與調試名的對應關系：?

debugger.extract_debug_model('before_debug.onnx', 'after_debug.onnx')

這步操作得到的?after_debug.onnx?模型的部分可視化結果如下：

我們可以使用下面的代碼運行這個模型：?

debugger.run_debug_model({'input':dummy_input.numpy()}, 'after_debug.onnx')

這樣，ONNX 模型的調試張量輸出會記錄在?debugger.onnx_value?里。?

總算，一切準備工作結束了。我們可以輕輕松松地用一行代碼輸出精度對齊的結果：?

debugger.print_debug_result()

這個函數大致會輸出以下內容：?

x_0 MSE: 8.465450562766819e-16  
x_1 MSE: 1.4122021817221354e-16  
x_2 MSE: 6.501743508551734e-17  
x_3 MSE: 1.7635199492054931e-16

這份輸出表明，在這一輪精度對齊測試中，所有模塊的精度誤差都很小。我們幾乎可以認為，ONNX 模型的運行結果等價于 PyTorch 模型的運行結果。?

如果有某些模塊的誤差比較大，我們可以深入子模塊，去加更多的 debug 節點，看看是哪一步、哪一個算子出現了問題。?