深度探索ONNX模型部署

0x0. 前言

這一節(jié)我將主要從盤點(diǎn)ONNX模型部署有哪些常見問題，以及針對(duì)這些問題提出一些解決方法，另外本文也會(huì)簡(jiǎn)單介紹一個(gè)可以快速用于ONNX模型推理驗(yàn)證的框架ONNXRuntime。如果你想用ONNX作為模型轉(zhuǎn)換和部署的工具，可以耐心看下去。今天要講到的ONNX模型部署碰到的問題大多來(lái)自于一些關(guān)于ONNX模型部署的文章以及自己使用ONNX進(jìn)行模型部署過程中的一些經(jīng)歷，有一定的實(shí)踐意義。

0x1. 導(dǎo)出ONNX

這里以Pytorch為例，來(lái)介紹一下要把Pytorch模型導(dǎo)出為ONNX模型需要注意的一些點(diǎn)。首先，Pytorch導(dǎo)出ONNX的代碼一般是這樣：

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

可以看到Pytorch提供了一個(gè)ONNX模型導(dǎo)出的專用接口，只需要配置好相關(guān)的模型和參數(shù)就可以完成自動(dòng)導(dǎo)出ONNX模型的操作了。代碼相關(guān)細(xì)節(jié)請(qǐng)自行查看，這里來(lái)列舉幾個(gè)導(dǎo)出ONNX模型中應(yīng)該注意的問題。

• 自定義OP問題

以2020年的YOLOV5為例，在模型的BackBone部分自定義了一個(gè)Focus OP，這個(gè)OP的代碼實(shí)現(xiàn)為：

class Focus(nn.Module):    # Focus wh information into c-space    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups        super(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))

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

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

綜上，自定義的OP在導(dǎo)出ONNX進(jìn)行部署時(shí)，除了考慮ONNX模型的執(zhí)行效率問題，還要考慮框架是否支持的問題。如果想快速迭代產(chǎn)品，建議盡量以一些經(jīng)典結(jié)構(gòu)為基礎(chǔ)，盡量少引入自定義OP。

• 后處理問題

如果我們要導(dǎo)出檢測(cè)網(wǎng)絡(luò)的ONNX模型進(jìn)行部署，就會(huì)碰到這個(gè)問題，后處理部分是否需要導(dǎo)入到ONNX模型？

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

檢測(cè)網(wǎng)絡(luò)比如YOLOV3的后處理就是NMS，代碼示例如https://github.com/ultralytics/yolov3/blob/master/utils/general.py#L325。當(dāng)我們完成檢測(cè)網(wǎng)絡(luò)的訓(xùn)練之后直接導(dǎo)出ONNX模型我們就會(huì)發(fā)現(xiàn)NMS這個(gè)實(shí)現(xiàn)也全部被導(dǎo)入了ONNX，如下圖所示：

這個(gè)結(jié)構(gòu)非常復(fù)雜，我們要在實(shí)際業(yè)務(wù)中去部署這個(gè)模型難度是很大的。另外，剛才我們提到ONNX模型只能記錄Pytorch中的Aten操作，對(duì)其它的一些邏輯運(yùn)算符比如if是無(wú)能為力的（意思是不能記錄if的多個(gè)子圖），而后處理過程中根據(jù)置信度閾值來(lái)篩選目標(biāo)框是常規(guī)操作。如果我們?cè)趯?dǎo)出ONNX模型時(shí)是隨機(jī)輸入或者沒有指定目標(biāo)的圖片就會(huì)導(dǎo)致這個(gè)ONNX記錄下來(lái)的DAG可能有缺失。最后，每個(gè)人實(shí)現(xiàn)后處理的方式可能都是不一樣的，這也增加了ONNX模型部署的難度。為了部署的友好性和降低轉(zhuǎn)換過程中的風(fēng)險(xiǎn)，后處理過程最好由讀者自己完成，我們只需要導(dǎo)出模型的Backbone和Neck部分為ONNX。

具體來(lái)說，我們只需要在Pytorch的代碼實(shí)現(xiàn)中屏蔽掉后處理部分然后導(dǎo)出ONNX模型即可。這也是目前使用ONNX部署檢測(cè)模型的通用方案。

所以，針對(duì)后處理問題，我們的結(jié)論就是在使用ONNX進(jìn)行部署時(shí)直接屏蔽后處理，將后處理單獨(dú)拿出來(lái)處理。

• 膠水OP問題。

在導(dǎo)出ONNX模型的過程中，經(jīng)常會(huì)帶來(lái)一些膠水OP，比如Gather, Shape等等。例如上節(jié)推文中介紹到當(dāng)執(zhí)行下面的Pytorch導(dǎo)出ONNX程序時(shí)，就會(huì)引入很多膠水OP。

import torch
class 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'])

導(dǎo)出的ONNX模型可視化如下：

這個(gè)時(shí)候的做法一般就是過一遍onnx-simplifer，可以去除這些膠水OP獲得一個(gè)簡(jiǎn)化后的模型。

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

0x2. ONNX or Caffe？

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

其實(shí)，我想說的應(yīng)該是Caffe是過去，而ONNX是將來(lái)。為什么要這樣說？

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

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

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

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

最后，越來(lái)越多的廠商推出的端側(cè)推理芯片開始支持ONNX，比如地平線的BPU，華為的Ascend310相關(guān)的工具鏈都開始接入ONNX，所以個(gè)人認(rèn)為ONNX是推理框架模型轉(zhuǎn)換的未來(lái)，不過仍需時(shí)間考驗(yàn)，畢竟誰(shuí)也不希望因?yàn)榭蚣躉P對(duì)齊的原因?qū)С鲆粋€(gè)超級(jí)復(fù)雜的ONNX模型，還是簡(jiǎn)化不了的那種，導(dǎo)致部署難度很大。

0x3. 一些典型坑點(diǎn)及解決辦法

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

• 預(yù)處理問題。

和后處理對(duì)應(yīng)的還有預(yù)處理問題，如果在Pytorch中使用下面的代碼導(dǎo)出ONNX模型。

import torch

class 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.std        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'])

我們先給這個(gè)ONNX模型過一遍onnx-simplifer，然后使用Netron可視化之后模型大概長(zhǎng)這樣：

如果我們要把這個(gè)模型放到NPU上部署，如果NPU芯片不支持Sub和Div的量化計(jì)算，那么這兩個(gè)操作會(huì)被回退到NPU上進(jìn)行計(jì)算，這顯然是不合理的，因?yàn)槲覀兛偸窍刖W(wǎng)絡(luò)在NPU上能一鏡到底，中間斷開必定會(huì)影響模型效率，所以這里的解決辦法就是把預(yù)處理放在基于nn.Module搭建模型的代碼之外，然后推理的時(shí)候先把預(yù)處理做掉即可。

• 框架OP實(shí)現(xiàn)不一致問題

當(dāng)從Mxnet轉(zhuǎn)換模型到ONNX時(shí)，如果模型是帶有PReLU OP的（在人臉識(shí)別網(wǎng)絡(luò)很常見），就是一個(gè)大坑了。主要有兩個(gè)問題，當(dāng)從mxnet轉(zhuǎn)為ONNX時(shí)，PReLU的slope參數(shù)維度可能會(huì)導(dǎo)致onnxruntime推理時(shí)失敗，報(bào)錯(cuò)大概長(zhǎng)這樣：

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

這個(gè)錯(cuò)誤產(chǎn)生的原因可能是MxNet的版本問題（https://github.com/apache/incubator-mxnet/issues/17821），這個(gè)時(shí)候的解決辦法就是：修改PRelu層的slope參數(shù)的shape，不僅包括type參數(shù)，對(duì)應(yīng)的slope值也要修改來(lái)和shape對(duì)應(yīng)。

核心代碼如下：

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])

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

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

所以，我們?cè)趏nnx模型中只需要按照這種方法將PReLU節(jié)點(diǎn)進(jìn)行等價(jià)替換就可以了。

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

• 其它問題

當(dāng)我們使用tf2onnx工具將TensorFlow模型轉(zhuǎn)為ONNX模型時(shí)，模型的輸入batch維度沒有被設(shè)置，我們需要自行添加。解決代碼如下：

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

當(dāng)我們基于ONNX和TensorRT部署風(fēng)格遷移模型，里面有Instance Norm OP的時(shí)候，可能會(huì)發(fā)現(xiàn)結(jié)果不準(zhǔn)確，這個(gè)問題在這里被提出：https://forums.developer.nvidia.com/t/inference-result-inaccurate-with-conv-and-instancenormalization-under-certain-conditions/111617。經(jīng)過debug發(fā)現(xiàn)這個(gè)問題出在這里：https://github.com/onnx/onnx-tensorrt/blob/5dca8737851118f6ab8a33ea1f7bcb7c9f06caf5/builtin_op_importers.cpp#L1557。

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

# 給ONNX模型中的目標(biāo)節(jié)點(diǎn)設(shè)置指定屬性# 調(diào)用方式為：set_node_attribute(in_node, "epsilon", 1e-5)# 其中in_node就是所有的instancenorm op。    def set_node_attribute(self, target_node, attr_name, attr_value):        flag = False        for attr in target_node.attribute:            if (attr.name == attr_name):                if attr.type == 1:                    attr.f = attr_value                elif attr.type == 2:                    attr.i = attr_value                elif attr.type == 3:                    attr.s = attr_value                elif attr.type == 4:                    attr.t = attr_value                elif 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_value                elif attr.type == 7:                    attr.ints[:] = attr_value                elif attr.type == 8:                    attr.strings[:] = attr_value                else:                    print("unsupported attribute data type with attribute name")                    return False                flag = True
        if not flag:            # attribute not in original node            print("Warning: you are appending a new attribute to the node!")            target_node.attribute.append(helper.make_attribute(attr_name, attr_value))            flag = True        return flag

當(dāng)我們使用了Pytorch里面的[]索引操作或者其它需要判斷的情況，ONNX模型會(huì)多出一些if OP，這個(gè)時(shí)候這個(gè)if OP的輸入已經(jīng)是一個(gè)確定的True，因?yàn)槲覀円呀?jīng)介紹過為False那部分的子圖會(huì)被丟掉。這個(gè)時(shí)候建議過一遍最新的onnx-simplifier或者手動(dòng)刪除所有的if OP，代碼實(shí)現(xiàn)如下：


# 通過op的類型獲取onnx模型的計(jì)算節(jié)點(diǎn)    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模型中的目標(biāo)節(jié)點(diǎn)    def remove_node(self, target_node):        '''            刪除只有一個(gè)輸入和輸出的節(jié)點(diǎn)        '''        node_input = target_node.input[0]        node_output = target_node.output[0]        # 將后繼節(jié)點(diǎn)的輸入設(shè)置為目標(biāo)節(jié)點(diǎn)的前置節(jié)點(diǎn)        for node in self.model.graph.node:            for i, n in enumerate(node.input):                if n == node_output:                    node.input[i] = node_input
        target_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，然后刪除這些節(jié)點(diǎn)。

0x4. ONNXRuntime介紹及用法

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

import numpy as npimport onnximport onnxruntime as ort
image = 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].nameoutput_name = sess.get_outputs()[0].name
output = sess.run([output_name], {input_name : image_data})prob = np.squeeze(output[0])print("predicting label:", np.argmax(prob))

這里展示了一個(gè)使用ONNXRuntime推理ResNet18網(wǎng)絡(luò)模型的例子，可以看到ONNXRuntime在推理一個(gè)ONNX模型時(shí)大概分為Session構(gòu)造，模型加載與初始化和運(yùn)行階段（和靜態(tài)圖框架類似）。ONNXRuntime框架是使用C++開發(fā)，同時(shí)使用Wapper技術(shù)封裝了Python接口易于用戶使用。

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

0x5. 調(diào)試工具

會(huì)逐漸補(bǔ)充一些解決ONNX模型出現(xiàn)的BUG或者修改，調(diào)試ONNX模型的代碼到這里：https://github.com/BBuf/onnx_learn 。這一節(jié)主要介紹幾個(gè)工具類函數(shù)結(jié)合ONNXRuntime來(lái)調(diào)試ONNX模型。

假設(shè)我們通過Pytorch導(dǎo)出了一個(gè)ONNX模型，在和Pytorch有相同輸入的情況下輸出結(jié)果卻不正確。這個(gè)時(shí)候我們要定位問題肯定需要獲取ONNX模型指定OP的特征值進(jìn)行對(duì)比，但是ONNX模型的輸出在導(dǎo)出模型的時(shí)候已經(jīng)固定了，這個(gè)時(shí)候應(yīng)該怎么做？

首先，我們需要通過名字獲取ONNX模型中的計(jì)算節(jié)點(diǎn)，實(shí)現(xiàn)如下：

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

然后把這個(gè)我們想看的節(jié)點(diǎn)擴(kuò)展到ONNX的輸出節(jié)點(diǎn)列表里面去，實(shí)現(xiàn)如下：

 # 將target_node添加到ONNX模型中作為輸出節(jié)點(diǎn)    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_type                for 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推理程序中的輸出節(jié)點(diǎn)為我們指定的節(jié)點(diǎn)就可以拿到指定節(jié)點(diǎn)的推理結(jié)果了，和Pytorch對(duì)比一下我們就可以知道是哪一層出錯(cuò)了。

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

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

查看指定節(jié)點(diǎn)的輸入節(jié)點(diǎn)的名字實(shí)現(xiàn)如下：

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

0x6. 總結(jié)

這篇文章從多個(gè)角度探索了ONNX，從ONNX的導(dǎo)出到ONNX和Caffe的對(duì)比，以及使用ONNX遭遇的困難以及一些解決辦法，另外還介紹了ONNXRuntime以及如何基于ONNXRuntime來(lái)調(diào)試ONNX模型等，后續(xù)會(huì)繼續(xù)結(jié)合ONNX做一些探索性工作。

參考資料：

• 資料1：https://zhuanlan.zhihu.com/p/128974102
• 資料2：https://zhuanlan.zhihu.com/p/165294876
• 資料3：https://zhuanlan.zhihu.com/p/212893519
• 資料4：https://blog.csdn.net/zsf10220208/article/details/107457820
• 資料5：https://github.com/bindog/onnx-surgery
• 資料6：https://zhuanlan.zhihu.com/p/346544539
• 資料7：https://github.com/daquexian/onnx-simplifier