ONNX再探

本文轉自：https://blog.csdn.net/just_sort/article/details/113802330

這篇文章從多個角度探索了ONNX，從ONNX的導出到ONNX和Caffe的對比，以及使用ONNX遭遇的困難以及一些解決辦法，另外還介紹了ONNXRuntime以及如何基于ONNXRuntime來調試ONNX模型等，后續也會繼續結合ONNX做一些探索性工作。

0x0. 前言

接著上篇文章，繼續探索ONNX。這一節我將主要從盤點ONNX模型部署有哪些常見問題，以及針對這些問題提出一些解決方法，另外本文也會簡單介紹一個可以快速用于ONNX模型推理驗證的框架ONNXRuntime。如果你想用ONNX作為模型轉換和部署的工具，可以耐心看下去。今天要講到的ONNX模型部署碰到的問題大多來自于一些關于ONNX模型部署的文章以及自己使用ONNX進行模型部署過程中的一些經歷，有一定的實踐意義。

0x1. 導出ONNX

這里以Pytorch為例，來介紹一下要把Pytorch模型導出為ONNX模型需要注意的一些點。首先，Pytorch導出ONNX的代碼一般是這樣：

import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = torch.load("test.pth") # pytorch模型加載
batch_size = 1  #批處理大小
input_shape = (3, 244, 224)   #輸入數據,改成自己的輸入shape# #set the model to inference mode
model.eval()x = torch.randn(batch_size, *input_shape)   # 生成張量
x = x.to(device)
export_onnx_file = "test.onnx"		# 目的ONNX文件名
torch.onnx.export(modelx,export_onnx_file,opset_version=10,do_constant_folding=True,	# 是否執行常量折疊優化input_names=["input"],	# 輸入名output_names=["output"],	# 輸出名dynamic_axes={"input":{0:"batch_size"},  # 批處理變量"output":{0:"batch_size"}})

可以看到Pytorch提供了一個ONNX模型導出的專用接口，只需要配置好相關的模型和參數就可以完成自動導出ONNX模型的操作了。代碼相關細節請自行查看，這里來列舉幾個導出ONNX模型中應該注意的問題。

自定義OP問題

以2020年的YOLOV5為例，在模型的BackBone部分自定義了一個Focus OP，這個OP的代碼實現為：

class Focus(nn.Module):# Focus wh information into c-spacedef __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groupssuper(Focus, self).__init__()self.conv = Conv(c1 * 4, c2, k, s, p, g, act)# self.contract = Contract(gain=2)def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))# return self.conv(self.contract(x))

這個操作就是一個stride slice然后再concat的操作，類似于PixelShuffle的逆向過程。下面是把YOLOv5模型導出ONNX模型之后Focus層的可視化結果。

在這里插入圖片描述

可以看到這個OP在使用Pytorch導出ONNX的過程中被拆成了很多更小的操作，這個時候Focus OP的問題就是推理的效率可能比較低并且拆成的小OP各個推理框架的支持程度不一致。要解決這種問題，要么直接在前向推理框架實現一個自定義的Focus OP，ncnn在實現yolov5的時候也是這樣做的：https://github.com/Tencent/ncnn/blob/master/examples/yolov5.cpp#L24。要么將這個OP使用其它的操作來近似代替，比如這里可以使用一個stride為2的卷積OP來代替Focus結構，注意代替之后有可能準確率會下降，需要做精度和部署友好性的平衡。

綜上，自定義的OP在導出ONNX進行部署時，除了考慮ONNX模型的執行效率問題，還要考慮框架是否支持的問題。如果想快速迭代產品，建議盡量以一些經典結構為基礎，盡量少引入自定義OP。

后處理的問題

如果我們要導出檢測網絡的ONNX模型進行部署，就會碰到這個問題，后處理部分是否需要導入到ONNX模型？

我們知道在使用Pytorch導出ONNX模型時，所有的Aten操作都會被ONNX記錄下來（具體記錄什么內容請參考文章開頭鏈接推文的介紹），成為一個DAG。然后ONNX會根據這個DAG的輸出節點來反推這個DAG中有哪些節點是有用的，這樣獲得的就是最終的ONNX模型。而后處理，比如非極大值抑制也是通過Aten操作拼起來的，所謂Aten操作就是Pytorch中的基礎算術單元比如加減乘除，所有的OP以及和Tensor相關的操作都基于Aten中的操作拼。

檢測網絡比如YOLOV3的后處理就是NMS，代碼示例如https://github.com/ultralytics/yolov3/blob/master/utils/general.py#L325。當我們完成檢測網絡的訓練之后直接導出ONNX模型我們就會發現NMS這個實現也全部被導入了ONNX，如下圖所示：

在這里插入圖片描述

這個結構非常復雜，我們要在實際業務中去部署這個模型難度是很大的。另外，剛才我們提到ONNX模型只能記錄Pytorch中的Aten操作，對其它的一些邏輯運算符比如 if 是無能為力的（意思是不能記錄if的多個子圖），而后處理過程中根據置信度閾值來篩選目標框是常規操作。如果我們在導出ONNX模型時是隨機輸入或者沒有指定目標的圖片就會導致這個ONNX記錄下來的DAG可能有缺失。最后，每個人實現后處理的方式可能都是不一樣的，這也增加了ONNX模型部署的難度。為了部署的友好性和降低轉換過程中的風險，后處理過程最好由讀者自己完成，我們只需要導出模型的Backbone和Neck部分為ONNX。

具體來說，我們只需要在Pytorch的代碼實現中屏蔽掉后處理部分然后導出ONNX模型即可。這也是目前使用ONNX部署檢測模型的通用方案。

所以，針對后處理問題，我們的結論就是在使用ONNX進行部署時直接屏蔽后處理，將后處理單獨拿出來處理。

膠水OP問題

在導出ONNX模型的過程中，經常會帶來一些膠水OP，比如Gather, Shape等等。例如上節推文中介紹到當執行下面的Pytorch導出ONNX程序時，就會引入很多膠水OP。

import torchclass JustReshape(torch.nn.Module):def __init__(self):super(JustReshape, self).__init__()def forward(self, x):return x.view((x.shape[0], x.shape[1], x.shape[3], x.shape[2]))net = JustReshape()
model_name = '../model/just_reshape.onnx'
dummy_input = torch.randn(2, 3, 4, 5)
torch.onnx.export(net, dummy_input, model_name, input_names=['input'], output_names=['output'])

導出的ONNX模型可視化如下：

在這里插入圖片描述

這個時候的做法一般就是過一遍onnx-simplifer，可以去除這些膠水OP獲得一個簡化后的模型。

在這里插入圖片描述

綜上，我們在導出ONNX模型的一般流程就是，去掉后處理，盡量不引入自定義OP，然后導出ONNX模型，并過一遍大老師的https://github.com/daquexian/onnx-simplifier，這樣就可以獲得一個精簡的易于部署的ONNX模型。從ONNX官方倉庫提供的模型來看，似乎微軟真的想用ONNX來統一所有框架的所有操作。但理想很豐滿，現實很骨干，各種訓練框架的數據排布，OP實現不一致，人為后處理不一致，各種推理框架支持度不一致，推理芯片SDK的OP支持度不一致都讓這個ONNX（萬能格式）遭遇了困難，所以在基于ONNX做一些部署業務的時候，也要有清晰的判斷并選取風險最小的方法。

0x2. ONNX or Caffe？

這個問題其實源于之前做模型轉換和基于TensorRT部署一些模型時候的思考。我們還是以Pytorch為例，要把Pytorch模型通過TensorRT部署到GPU上，一般就是Pytorch->Caffe->TensorRT以及Pytorch->ONNX->TensorRT（當然Pytorch也是支持直接轉換到TensorRT，這里不關心）。那么這里就有一個問題，我們選擇哪一條路比較好？

其實，我想說的應該是Caffe是過去，而ONNX是將來。為什么要這樣說？

首先很多國產推理芯片比如海思NNIE，高通SNPE它們首先支持的都是Caffe這種模型格式，這可能是因為年代的原因，也有可能是因為這些推理SDK實現的時候OP都非常粗粒度。比如它對卷積做定制的優化，有NC4HW4，有Im2Col+gemm，有Winograd等等非常多方法，后面還考慮到量化，半精度等等，然后通過給它喂Caffe模型它就知道要對這個網絡里面對應的卷積層進行硬件加速了。所以這些芯片支持的網絡是有限的，比如我們要在Hisi35xx中部署一個含有upsample層的Pytorch模型是比較麻煩的，可能不太聰明的工程說我們要把這個模型回退給訓練人員改成支持的上采樣方式進行訓練，而聰明的工程師可能說直接把upsample的參數填到反卷積層的參數就可以了。無論是哪種方式都是比較麻煩的，所以Caffe的缺點就是靈活度太差。其實基于Caffe進行部署的方式仍然在工業界發力，ONNX是趨勢，但是ONNX現在還沒有完全取代Caffe。

接下來，我們要再提一下上面那個 if 的事情了，假設現在有一個新的SOTA模型被提出，這個模型有一個自定義的OP，作者是用Pytorch的Aten操作拼的，邏輯大概是這樣：

result = check()
if result == 0:result = algorithm1(result)
else:result = algorithm2(result)
return result

然后考慮將這個模型導出ONNX或者轉換為Caffe，如果是Caffe的話我們需要去實現這個自定義的OP，并將其注冊到Caffe的OP管理文件中，雖然這里比較繁瑣，但是我們可以將if操作隱藏在這個大的OP內部，這個 if 操作可以保留下來。而如果我們通過導出ONNX模型的方式 if 子圖只能保留一部分，要么保留algorithm1，要么保留algorithm2對應的子圖，這種情況ONNX似乎就沒辦法處理了。這個時候要么保存兩個ONNX模型，要么修改算法邏輯繞過這個問題。從這里引申一下，如果我們碰到有遞歸關系的網絡，基于ONNX應當怎么部署？ONNX還有一個缺點就是OP的細粒度太細，執行效率低，不過ONNX已經推出了多種化方法可以將OP的細粒度變粗，提高模型執行效率。目前在眾多經典算法上，ONNX已經支持得非常好了。

最后，越來越多的廠商推出的端側推理芯片開始支持ONNX，比如地平線的BPU，華為的Ascend310相關的工具鏈都開始接入ONNX，所以個人認為ONNX是推理框架模型轉換的未來，不過仍需時間考驗，畢竟誰也不希望因為框架OP對齊的原因導出一個超級復雜的ONNX模型，還是簡化不了的那種，導致部署難度很大。

0x3. 一些典型的坑點及解決辦法

第一節已經提到，將我們的ONNX模型過一遍onnx-simplifer之后就可以去掉膠水OP并將一些細粒度的OP進行op fuse成粗粒度的OP，并解決一部分由于Pytorch和ONNX OP實現方式不一致而導致模型變復雜的問題。除了這些問題，本節再列舉一些ONNX模型部署中容易碰到的坑點，并嘗試給出一些解決辦法。

預處理問題

和后處理對應的還有預處理問題，如果在Pytorch中使用下面的代碼導出ONNX模型。

import torchclass JustReshape(torch.nn.Module):def __init__(self):super(JustReshape, self).__init__()self.mean = torch.randn(2, 3, 4, 5)self.std = torch.randn(2, 3, 4, 5)def forward(self, x):x = (x - self.mean) / self.stdreturn x.view((x.shape[0], x.shape[1], x.shape[3], x.shape[2]))net = JustReshape()
model_name = '../model/just_reshape.onnx'
dummy_input = torch.randn(2, 3, 4, 5)
torch.onnx.export(net, dummy_input, model_name, input_names=['input'], output_names=['output'])

我們先給這個ONNX模型過一遍onnx-simplifer，然后使用Netron可視化之后模型大概長這樣：
在這里插入圖片描述

如果我們要把這個模型放到NPU上部署，如果NPU芯片不支持Sub和Div的量化計算，那么這兩個操作會被回退到NPU上進行計算，這顯然是不合理的，因為我們總是想網絡在NPU上能一鏡到底，中間斷開必定會影響模型效率，所以這里的解決辦法就是把預處理放在基于nn.Module搭建模型的代碼之外，然后推理的時候先把預處理做掉即可。

框架OP實現的問題

當從Mxnet轉換模型到ONNX時，如果模型是帶有PReLU OP的（在人臉識別網絡很常見），就是一個大坑了。主要有兩個問題，當從mxnet轉為ONNX時，PReLU的slope參數維度可能會導致onnxruntime推理時失敗，報錯大概長這樣：

2）[ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running PRelu node. Name:'conv_1_relu'...... Attempting to broadcast an axis by a dimension other than 1. 56 by 64

這個錯誤產生的原因可能是MxNet的版本問題（https://github.com/apache/incubator-mxnet/issues/17821），這個時候的解決辦法就是：修改PRelu層的slope參數的shape，不僅包括type參數，對應的slope值也要修改來和shape對應。

核心代碼如下：

graph.input.remove(input_map[input_name])
new_nv = helper.make_tensor_value_info(input_name, TensorProto.FLOAT, [input_dim_val,1,1])
graph.input.extend([new_nv])

想了解更加詳細的信息可以參考問后的資料2和資料3。

這個問題其實算是個小問題，我們自己在ONNX模型上fix一下即可。下一個問題就是如果我們將處理好之后的ONNX通過TensorRT進行部署的話，我們會發現TensorRT不支持PReLU這個OP，這個時候解決辦法要么是TensorRT自定義PReLU插件，但是這種方法會打破TensorRT中conv+bn+relu的op fusion，速度會變慢，并且如果要做量化部署似乎是不可行的。所以這個時候一般會采用另外一種解決辦法，使用relu和scale op來組合成PReLU，如下圖所示：

在這里插入圖片描述

所以，我們在onnx模型中只需要按照這種方法將PReLU節點進行等價替換就可以了。

這個地方以PReLU列舉了一個框架OP實現不一致的問題，比如大老師最新文章也介紹的就是squeeze OP在Pytorch和ONNX實現時的不一致導致ONNX模型變得很復雜，這種問題感覺是基于ONNX支持模型部署時的常見問題，雖然onnx-simplifier已經解決了一些問題，但也不能夠完全解決。

其他問題

當我們使用tf2onnx工具將TensorFlow模型轉為ONNX模型時，模型的輸入batch維度沒有被設置，我們需要自行添加。解決代碼如下：

# 為onnx模型增加batch維度def set_model_input_batch(self, index=0, name=None, batch_size=4):model_input = Noneif name is not None:for ipt in self.model.graph.input:if ipt.name == name:model_input = iptelse:model_input = self.model.graph.input[index]if model_input:tensor_dim = model_input.type.tensor_type.shape.dimtensor_dim[0].ClearField("dim_param")tensor_dim[0].dim_value = batch_sizeelse:print('get model input failed, check index or name')

當我們基于ONNX和TensorRT部署風格遷移模型，里面有Instance Norm OP的時候，可能會發現結果不準確，這個問題在這里被提出：https://forums.developer.nvidia.com/t/inference-result-inaccurate-with-conv-and-instancenormalization-under-certain-conditions/111617。經過debug發現這個問題出在這里：https://github.com/onnx/onnx-tensorrt/blob/5dca8737851118f6ab8a33ea1f7bcb7c9f06caf5/builtin_op_importers.cpp#L1557。

在這里插入圖片描述

問題比較明顯了，instancenorm op里面的eps只支持>=1e-4的，所以要么注釋掉這個限制條件，要么直接在ONNX模型中修改instancenorm op的eps屬性，代碼實現如下：

# 給ONNX模型中的目標節點設置指定屬性
# 調用方式為：set_node_attribute(in_node, "epsilon", 1e-5)
# 其中in_node就是所有的instancenorm op。def set_node_attribute(self, target_node, attr_name, attr_value):flag = Falsefor attr in target_node.attribute:if (attr.name == attr_name):if attr.type == 1:attr.f = attr_valueelif attr.type == 2:attr.i = attr_valueelif attr.type == 3:attr.s = attr_valueelif attr.type == 4:attr.t = attr_valueelif attr.type == 5:attr.g = attr_value# NOTE: For repeated composite types, we should use something like# del attr.xxx[:]# attr.xxx.extend([n1, n2, n3])elif attr.type == 6:attr.floats[:] = attr_valueelif attr.type == 7:attr.ints[:] = attr_valueelif attr.type == 8:attr.strings[:] = attr_valueelse:print("unsupported attribute data type with attribute name")return Falseflag = Trueif not flag:# attribute not in original nodeprint("Warning: you are appending a new attribute to the node!")target_node.attribute.append(helper.make_attribute(attr_name, attr_value))flag = Truereturn flag

當我們使用了Pytorch里面的[]索引操作或者其它需要判斷的情況，ONNX模型會多出一些if OP，這個時候這個if OP的輸入已經是一個確定的True，因為我們已經介紹過為False那部分的子圖會被丟掉。這個時候建議過一遍最新的onnx-simplifier或者手動刪除所有的if OP，代碼實現如下：

# 通過op的類型獲取onnx模型的計算節點def get_nodes_by_optype(self, typename):nodes = []for node in self.model.graph.node:if node.op_type == typename:nodes.append(node)return nodes
# 移除ONNX模型中的目標節點def remove_node(self, target_node):'''刪除只有一個輸入和輸出的節點'''node_input = target_node.input[0]node_output = target_node.output[0]# 將后繼節點的輸入設置為目標節點的前置節點for node in self.model.graph.node:for i, n in enumerate(node.input):if n == node_output:node.input[i] = node_inputtarget_names = set(target_node.input) & set([weight.name for weight in self.model.graph.initializer])self.remove_weights(target_names)target_names.add(node_output)self.remove_inputs(target_names)self.remove_value_infos(target_names)self.model.graph.node.remove(target_node)

具體順序就是先獲取所有if類型的OP，然后刪除這些節點。

0x4. ONNXRuntime介紹及用法

ONNXRuntime是微軟推出的一個推理框架，似乎最新版都支持訓練功能了，用戶可以非常方便的運行ONNX模型。ONNXRuntime支持多種運行后端包括CPU，GPU，TensorRT，DML等。ONNXRuntime是專為ONNX打造的框架，雖然我們大多數人把ONNX只是當成工具人，但微軟可不這樣想，ONNX統一所有框架的IR表示應該是終極理想吧。從使用者的角度我們簡單分析一下ONNXRuntime即可。

import numpy as np
import onnx
import onnxruntime as ortimage = cv2.imread("image.jpg")
image = np.expand_dims(image, axis=0)onnx_model = onnx.load_model("resnet18.onnx")
sess = ort.InferenceSession(onnx_model.SerializeToString())
sess.set_providers(['CPUExecutionProvider'])
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].nameoutput = sess.run([output_name], {input_name : image_data})
prob = np.squeeze(output[0])
print("predicting label:", np.argmax(prob))

這里展示了一個使用ONNXRuntime推理ResNet18網絡模型的例子，可以看到ONNXRuntime在推理一個ONNX模型時大概分為Session構造，模型加載與初始化和運行階段（和靜態圖框架類似）。ONNXRuntime框架是使用C++開發，同時使用Wapper技術封裝了Python接口易于用戶使用。

從使用者的角度來說，知道怎么用就可以了，如果要了解框架內部的知識請移步源碼（https://github.com/microsoft/onnxruntime）和參考資料6。

0x5. 調試工具

會逐漸補充一些解決ONNX模型出現的BUG或者修改，調試ONNX模型的代碼到這里：https://github.com/BBuf/onnx_learn 。這一節主要介紹幾個工具類函數結合ONNXRuntime來調試ONNX模型。

假設我們通過Pytorch導出了一個ONNX模型，在和Pytorch有相同輸入的情況下輸出結果卻不正確。這個時候我們要定位問題肯定需要獲取ONNX模型指定OP的特征值進行對比，但是ONNX模型的輸出在導出模型的時候已經固定了，這個時候應該怎么做？

首先，我們需要通過名字獲取ONNX模型中的計算節點，實現如下：

# 通過名字獲取onnx模型中的計算節點def get_node_by_name(self, name):for node in self.model.graph.node:if node.name == name:return node

然后把這個我們想看的節點擴展到ONNX的輸出節點列表里面去，實現如下：

 # 將target_node添加到ONNX模型中作為輸出節點def add_extra_output(self, target_node, output_name):target_output = target_node.output[0]extra_shape = []for vi in self.model.graph.value_info:if vi.name == target_output:extra_elem_type = vi.type.tensor_type.elem_typefor s in vi.type.tensor_type.shape.dim:extra_shape.append(s.dim_value)extra_output = helper.make_tensor_value_info(output_name,extra_elem_type,extra_shape)identity_node = helper.make_node('Identity', inputs=[target_output], outputs=[output_name], name=output_name)self.model.graph.node.append(identity_node)self.model.graph.output.append(extra_output)

然后修改一下onnxruntime推理程序中的輸出節點為我們指定的節點就可以拿到指定節點的推理結果了，和Pytorch對比一下我們就可以知道是哪一層出錯了。

這里介紹的是如何查看ONNX在確定輸入的情況下如何拿到推理結果，如果我們想要獲取ONNX模型中某個節點的信息又可以怎么做呢？這個就結合上一次推文講的ONNX的結構來看就比較容易了。例如查看某個指定節點的屬性代碼實現如下：

def show_node_attributes(node):print("="*10, "attributes of node: ", node.name, "="*10)for attr in node.attribute:print(attr.name)print("="*60)

查看指定節點的輸入節點的名字實現如下：

def show_node_inputs(node):# Generally, the first input is the truely input# and the rest input is weight initializerprint("="*10, "inputs of node: ", node.name, "="*10)for input_name in node.input:print(input_name)  # type of input_name is strprint("="*60)