實(shí)操教程｜Pytorch轉(zhuǎn)ONNX詳解

來(lái)源丨h(huán)ttps://zhuanlan.zhihu.com/p/272767300編輯丨極市平臺(tái)

之前幾個(gè)月參與了OpenMMlab的模型轉(zhuǎn)ONNX的工作（github account: drcut），主要目標(biāo)是支持OpenMMLab的一些模型從Pytorch到ONNX的轉(zhuǎn)換。這幾個(gè)月雖然沒(méi)做出什么成果，但是踩了很多坑，在這里記錄下來(lái)，希望可以幫助其他人。

這篇是第一部分，理論篇，主要介紹了和代碼無(wú)關(guān)的一些宏觀問(wèn)題。再接下來(lái)我會(huì)專門寫一篇實(shí)戰(zhàn)篇，針對(duì)OpenMMlab中一些具體代碼做分析，說(shuō)明Pytorch轉(zhuǎn)化ONNX過(guò)程中的一些代碼上的技巧和注意事項(xiàng)。

（1）Pytorch轉(zhuǎn)ONNX的意義

一般來(lái)說(shuō)轉(zhuǎn)ONNX只是一個(gè)手段，在之后得到ONNX模型后還需要再將它做轉(zhuǎn)換，比如轉(zhuǎn)換到TensorRT上完成部署，或者有的人多加一步，從ONNX先轉(zhuǎn)換到caffe，再?gòu)腸affe到tensorRT。原因是Caffe對(duì)tensorRT更為友好，這里關(guān)于友好的定義后面會(huì)談。

因此在轉(zhuǎn)ONNX工作開(kāi)展之前，首先必須明確目標(biāo)后端。ONNX只是一個(gè)格式，就和json一樣。只要你滿足一定的規(guī)則，都算是合法的，因此單純從Pytorch轉(zhuǎn)成一個(gè)ONNX文件很簡(jiǎn)單。但是不同后端設(shè)備接受的onnx是不一樣的，因此這才是坑的來(lái)源。

Pytorch自帶的torch.onnx.export轉(zhuǎn)換得到的ONNX，ONNXRuntime需要的ONNX，TensorRT需要的ONNX都是不同的。

這里面舉一個(gè)最簡(jiǎn)單的Maxpool的例：

Maxunpool可以被看作Maxpool的逆運(yùn)算，咱們先來(lái)看一個(gè)Maxpool的例子，假設(shè)有如下一個(gè)C*H*W的tensor（shape[2, 3, 3])，其中每個(gè)channel的二維矩陣都是一樣的，如下所示

在這種情況下，如果我們?cè)赑ytorch對(duì)它調(diào)用MaxPool（kernel_size=2, stride=1,pad=0）

那么會(huì)得到兩個(gè)輸出，第一個(gè)輸出是Maxpool之后的值：

另一個(gè)是Maxpool的Idx，即每個(gè)輸出對(duì)應(yīng)原來(lái)的哪個(gè)輸入，這樣做反向傳播的時(shí)候就可以直接把輸出的梯度傳給對(duì)應(yīng)的輸入：

細(xì)心的同學(xué)會(huì)發(fā)現(xiàn)其實(shí)Maxpool的Idx還可以有另一種寫法：

  ，

即每個(gè)channel的idx放到一起，并不是每個(gè)channel單獨(dú)從0開(kāi)始。這兩種寫法都沒(méi)什么問(wèn)題，畢竟只要反向傳播的時(shí)候一致就可以。

但是當(dāng)我在支持OpenMMEditing的時(shí)候，會(huì)涉及到Maxunpool，即Maxpool的逆運(yùn)算：輸入MaxpoolId和Maxpool的輸出，得到Maxpool的輸入。

Pytorch的MaxUnpool實(shí)現(xiàn)是接收每個(gè)channel都從0開(kāi)始的Idx格式，而Onnxruntime則相反。因此如果你希望用Onnxruntime跑一樣的結(jié)果，那么必須對(duì)輸入的Idx（即和Pytorch一樣的輸入）做額外的處理才可以。換言之，Pytorch轉(zhuǎn)出來(lái)的神經(jīng)網(wǎng)絡(luò)圖和ONNXRuntime需要的神經(jīng)網(wǎng)絡(luò)圖是不一樣的。

（2）ONNX與Caffe

主流的模型部署有兩種路徑，以TensorRT為例，一種是Pytorch->ONNX->TensorRT，另一種是Pytorch->Caffe->TensorRT。個(gè)人認(rèn)為目前后者更為成熟，這主要是ONNX，Caffe和TensorRT的性質(zhì)共同決定的

上面的表列了ONNX和Caffe的幾點(diǎn)區(qū)別，其中最重要的區(qū)別就是op的粒度。舉個(gè)例子，如果對(duì)Bert的Attention層做轉(zhuǎn)換，ONNX會(huì)把它變成MatMul，Scale，SoftMax的組合，而Caffe可能會(huì)直接生成一個(gè)叫做Multi-Head Attention的層，同時(shí)告訴CUDA工程師：“你去給我寫一個(gè)大kernel“（很懷疑發(fā)展到最后會(huì)不會(huì)把ResNet50都變成一個(gè)層。。。）

因此如果某天一個(gè)研究員提了一個(gè)新的State-of-the-art的op，很可能它直接就可以被轉(zhuǎn)換成ONNX（如果這個(gè)op在Pytorch的實(shí)現(xiàn)全都是用Aten的庫(kù)拼接的），但是對(duì)于Caffe的工程師，需要重新寫一個(gè)kernel。
細(xì)粒度op的好處就是非常靈活，壞處就是速度會(huì)比較慢。這幾年有很多工作都是在做op fushion（比如把卷積和它后面的relu合到一起算），XLA和TVM都有很多工作投入到了op fushion，也就是把小op拼成大op。
TensorRT是NVIDIA推出的部署框架，自然性能是首要考量的，因此他們的layer粒度都很粗。在這種情況下把Caffe轉(zhuǎn)換過(guò)去有天然的優(yōu)勢(shì)。
除此之外粗粒度也可以解決分支的問(wèn)題。TensorRT眼里的神經(jīng)網(wǎng)絡(luò)就是一個(gè)單純的DAG：給定固定shape的輸入，執(zhí)行相同的運(yùn)算，得到固定shape的輸出。
**目前TensorRT的一個(gè)發(fā)展方向是支持dynamic shape，但是還很不成熟。

tensor i = funcA();
if(i==0)
  j = funcB(i);
else
  j = funcC(i);
funcD(j);

對(duì)于上面的網(wǎng)絡(luò)，假設(shè)funcA,funcB,funcC和funcD都是onnx支持的細(xì)粒度算子，那么ONNX就會(huì)面臨一個(gè)困難，它轉(zhuǎn)換得到的DAG要么長(zhǎng)這樣：funcA->funcB->funcD，要么funcA->funcC->funcD。但是無(wú)論哪種肯定都是有問(wèn)題的。
而Caffe可以用粗粒度繞開(kāi)這個(gè)問(wèn)題

tensor i = funcA();
coarse_func(tensor i) {
  if(i==0) return funcB(i);
  else return funcC(i);
}
funcD(coarse_func(i))

因此它得到的DAG是：funcA->coarse_func->funcD
當(dāng)然，Caffe的代價(jià)就是苦逼的HPC工程師就要手寫一個(gè)coarse_func kernel。。。（希望Deep Learning Compiler可以早日解放HPC工程師）
(3)Pytorch本身的局限
熟悉深度學(xué)習(xí)框架的同學(xué)都知道，Pytorch之所以可以在tensorflow已經(jīng)占據(jù)主流的情況下橫空出世，成功搶占半壁江山，主要的原因是它很靈活。舉個(gè)不恰當(dāng)?shù)睦?，tensorflow就像是C++，而Pytorch就是Python。
tensorflow會(huì)把整個(gè)神經(jīng)網(wǎng)絡(luò)在運(yùn)行前做一次編譯，生成一個(gè)DAG（有向無(wú)環(huán)圖），然后再去跑這張圖。Pytorch則相反，屬于走一步看一步，直到運(yùn)行到這個(gè)節(jié)點(diǎn)算出結(jié)果，才知道下一個(gè)節(jié)點(diǎn)該算啥。
ONNX其實(shí)就是把上層深度學(xué)習(xí)框架中的網(wǎng)絡(luò)模型轉(zhuǎn)換成一張圖，因?yàn)閠ensorflow本身就有一張圖，因此只需要直接把這張圖拿到手，修修補(bǔ)補(bǔ)就可以。
但是對(duì)于Pytorch，沒(méi)有任何圖的概念，因此如果想完成Pytorch到ONNX的轉(zhuǎn)換，就需要讓ONNX再旁邊拿個(gè)小本子，然后跑一遍Pytorch，跑到什么就把什么記下來(lái)，把記錄的結(jié)果抽象成一張圖。因此Pytorch轉(zhuǎn)ONNX有兩個(gè)天然的局限。
1. 轉(zhuǎn)換的結(jié)果只對(duì)特定的輸入。如果換一個(gè)輸入導(dǎo)致網(wǎng)絡(luò)結(jié)構(gòu)發(fā)生了變化，ONNX是無(wú)法察覺(jué)的（最常見(jiàn)的情況是如果網(wǎng)絡(luò)中有if語(yǔ)句，這次的輸入走了if的話，ONNX就只會(huì)生成if對(duì)應(yīng)的圖，把else里面全部的信息都丟掉）。
2. 需要比較多的計(jì)算量，因?yàn)樾枰娴墩鏄尩呐芤槐樯窠?jīng)網(wǎng)絡(luò)。
PS：針對(duì)于以上的兩個(gè)局限，我的本科畢設(shè)論文提出了一種解決方案，就是通過(guò)編譯器里面的詞法分析，語(yǔ)法分析直接掃描Pytorch或者tensorflow的源代碼得到圖結(jié)構(gòu)，這樣可以輕量級(jí)的完成模型到ONNX的轉(zhuǎn)換，同時(shí)也可以得到分支判斷等信息，這里放一個(gè)github鏈接(https://github.com/drcut/NN_transform)，希望大家多多支持
*目前Pytorch官方希望通過(guò)用TorchScript的方式解決分支語(yǔ)句的問(wèn)題，但據(jù)我所知還不是很成熟。

本文僅做學(xué)術(shù)分享，如有侵權(quán)，請(qǐng)聯(lián)系刪文。