Pytorch轉(zhuǎn)ONNX詳解
點擊上方“AI算法與圖像處理”，選擇加"星標"或“置頂”
重磅干貨，第一時間送達
導(dǎo)讀
?本文作者總結(jié)了自己參與Pytorch到ONNX的模型轉(zhuǎn)換轉(zhuǎn)換工作中的經(jīng)驗，主要介紹了該轉(zhuǎn)換工作的意義，模型部署的路徑以及Pytorch本身的局限。?
之前幾個月參與了OpenMMlab的模型轉(zhuǎn)ONNX的工作（github account: drcut），主要目標是支持OpenMMLab的一些模型從Pytorch到ONNX的轉(zhuǎn)換。這幾個月雖然沒做出什么成果，但是踩了很多坑，在這里記錄下來，希望可以幫助其他人。
這篇是第一部分，理論篇，主要介紹了和代碼無關(guān)的一些宏觀問題。再接下來我會專門寫一篇實戰(zhàn)篇，針對OpenMMlab中一些具體代碼做分析，說明Pytorch轉(zhuǎn)化ONNX過程中的一些代碼上的技巧和注意事項。
（1）Pytorch轉(zhuǎn)ONNX的意義
一般來說轉(zhuǎn)ONNX只是一個手段，在之后得到ONNX模型后還需要再將它做轉(zhuǎn)換，比如轉(zhuǎn)換到TensorRT上完成部署，或者有的人多加一步，從ONNX先轉(zhuǎn)換到caffe，再從caffe到tensorRT。原因是Caffe對tensorRT更為友好，這里關(guān)于友好的定義后面會談。
因此在轉(zhuǎn)ONNX工作開展之前，首先必須明確目標后端。ONNX只是一個格式，就和json一樣。只要你滿足一定的規(guī)則，都算是合法的，因此單純從Pytorch轉(zhuǎn)成一個ONNX文件很簡單。但是不同后端設(shè)備接受的onnx是不一樣的，因此這才是坑的來源。
Pytorch自帶的torch.onnx.export轉(zhuǎn)換得到的ONNX，ONNXRuntime需要的ONNX，TensorRT需要的ONNX都是不同的。
這里面舉一個最簡單的Maxpool的例：
Maxunpool可以被看作Maxpool的逆運算，咱們先來看一個Maxpool的例子，假設(shè)有如下一個C*H*W的tensor（shape[2, 3, 3])，其中每個channel的二維矩陣都是一樣的，如下所示
?
在這種情況下，如果我們在Pytorch對它調(diào)用MaxPool（kernel_size=2, stride=1,pad=0）
那么會得到兩個輸出，第一個輸出是Maxpool之后的值：
?
另一個是Maxpool的Idx，即每個輸出對應(yīng)原來的哪個輸入，這樣做反向傳播的時候就可以直接把輸出的梯度傳給對應(yīng)的輸入：
?
細心的同學(xué)會發(fā)現(xiàn)其實Maxpool的Idx還可以有另一種寫法：
?
?，
即每個channel的idx放到一起，并不是每個channel單獨從0開始。這兩種寫法都沒什么問題，畢竟只要反向傳播的時候一致就可以。
但是當(dāng)我在支持OpenMMEditing的時候，會涉及到Maxunpool，即Maxpool的逆運算：輸入MaxpoolId和Maxpool的輸出，得到Maxpool的輸入。
Pytorch的MaxUnpool實現(xiàn)是接收每個channel都從0開始的Idx格式，而Onnxruntime則相反。因此如果你希望用Onnxruntime跑一樣的結(jié)果，那么必須對輸入的Idx（即和Pytorch一樣的輸入）做額外的處理才可以。換言之，Pytorch轉(zhuǎn)出來的神經(jīng)網(wǎng)絡(luò)圖和ONNXRuntime需要的神經(jīng)網(wǎng)絡(luò)圖是不一樣的。
（2）ONNX與Caffe
主流的模型部署有兩種路徑，以TensorRT為例，一種是Pytorch->ONNX->TensorRT，另一種是Pytorch->Caffe->TensorRT。個人認為目前后者更為成熟，這主要是ONNX，Caffe和TensorRT的性質(zhì)共同決定的
上面的表列了ONNX和Caffe的幾點區(qū)別，其中最重要的區(qū)別就是op的粒度。舉個例子，如果對Bert的Attention層做轉(zhuǎn)換，ONNX會把它變成MatMul，Scale，SoftMax的組合，而Caffe可能會直接生成一個叫做Multi-Head Attention的層，同時告訴CUDA工程師：“你去給我寫一個大kernel“（很懷疑發(fā)展到最后會不會把ResNet50都變成一個層。。。）
tensor i = funcA();
if(i==0)
j = funcB(i);
else
j = funcC(i);
funcD(j);
tensor i = funcA();
coarse_func(tensor i) {
if(i==0) return funcB(i);
else return funcC(i);
}
funcD(coarse_func(i))
下載1：何愷明頂會分享
在「AI算法與圖像處理」公眾號后臺回復(fù)：何愷明，即可下載。總共有6份PDF，涉及 ResNet、Mask RCNN等經(jīng)典工作的總結(jié)分析
下載2：leetcode?開源書
在「AI算法與圖像處理」公眾號后臺回復(fù)：leetcode，即可下載。每題都 runtime beats 100% 的開源好書，你值得擁有！
下載3 CVPR2020 在「AI算法與圖像處理」公眾號后臺回復(fù)：CVPR2020，即可下載1467篇CVPR?2020論文 個人微信（如果沒有備注不拉群！） 請注明：地區(qū)+學(xué)校/企業(yè)+研究方向+昵稱
覺得不錯就點亮在看吧
