鑽石舞台

在深度學習模型部署時，從pytorch轉換onnx的過程中，踩了一些坑。本文總結了這些踩坑記錄，希望可以幫助其他人。

首先，簡單說明一下pytorch轉onnx的意義。在pytorch訓練出一個深度學習模型後，需要在TensorRT或者openvino部署，這時需要先把Pytorch模型轉換到onnx模型之後再做其它轉換。因此，在使用pytorch訓練深度學習模型完成後，在TensorRT或者openvino或者opencv和onnxruntime部署時，pytorch模型轉onnx這一步是必不可少的。接下來通過幾個實例程序，介紹pytorch轉換onnx的過程中遇到的坑。

起初，我在微信公眾號里看到一篇文章《使用Python和YOLO檢測車牌》。文中展示的檢測結果如下，其實這種檢測結果並不是一個優良的結果，可以看到檢測框裡的車牌是傾斜的，如果要識別車牌里的文字，那麼傾斜的車牌會嚴重影響車牌識別結果的。

對於車牌識別這種場景，在做車牌檢測時，一種優良的檢測結果應該是這樣的，如下圖所示。

在輸出車牌檢測框的同時輸出檢測到的車牌的4個角點。有了這4個角點之後，對車牌做透視變換，這時的車牌就是水平放置的，最後做車牌識別，這樣就做成了一個車牌識別系統，在這個系統里包含車牌檢測，車牌矯正，車牌識別三個模塊。車牌檢測模塊使用retinaface，原始的retinaface是做人臉檢測的，它能輸出人臉檢測矩形框和人臉5個關鍵點。考慮到車牌只有4個點，於是修改retinaface的網絡結構使其輸出4個關鍵點，然後在車牌數據集訓練，訓練完成後，以一幅圖片上做目標檢測的結果如上圖所示。車牌矯正模塊使用了傳統圖像處理方法，關鍵函數是opencv里的getPerspectiveTransform和warpPerspective。車牌識別模塊使用Intel公司提出的LPRNet。

整套程序是基於pytorch框架運行的，我把這套程序發布在github上，地址是 https://github.com/hpc203/license-plate-detect-recoginition-pytorch

接下來我就嘗試把pytorch模型轉換到onnx文件，然後使用opencv做車牌檢測與識別。然而在轉換完成onnx文件後，使用opencv讀取onnx文件遇到了一些坑，我在網上搜索，也沒有找到解決辦法。

轉換過程分兩步，首先是轉換車牌檢測retinaface到onnx文件，這一步倒是很順利，轉換沒有出錯，並且使用opencv讀取onnx文件做前向推理的輸出結果也是正確的。第二步轉換車牌識別LPRNet到onnx文件，由於Pytorch自帶torch.onnx.export轉換得到的ONNX，因此轉換的代碼很簡單，在生成onnx文件後，opencv讀取onnx文件出現了模型其妙的錯誤。程序運行的結果截圖如下

從打印結果看，torch.onnx.export生成onnx文件時沒有問題的，但是在cv2.dnn.readNet這一步出現異常導致程序中斷，並且打印出的異常信息是一連串的數字，去百度搜索也麼找到解決辦法。觀察LPRNet的網絡結構，發現在LPRNet里定義了3維池化層，代碼截圖如下

可以看到依然出錯，這說明opencv的深度學習模塊里不支持3維池化。不過，對比3維池化和2維池化的前向計算原理可以發現，3維池化其實等價於2個2維池化。程序實例如下

程序最後最後運行結果打印信息是相等。從這裡就可以看出opencv里的深度學習模塊並不支持3維池化的前向計算，這期待後續新版本的opencv里能添加3維池化的計算。這時在LPRNet網絡結構定義文件里修改3維池化層，重新生成onnx文件，opencv讀取onnx文件執行前向計算後依然出錯，運行結果如下。

於是繼續觀察LPRNet的網絡結構，在forward函數裡看到有求平均值的操作，代碼截圖如下所示

注意到第一個torch.mean函數裡沒有聲明在哪個維度求平均值，這說明它是對一個4維四維張量的整體求平均值，這時候從一個4維空間搜索成一個點，也就是一個標量數值。但是在pytorch里，對一個張量求平均值後依然是一個張量，只不過它的維度shape是空的，示例代碼如下。這時如果想要訪問平均值，需要加上.item()，這個是需要注意的一個pytorch知識點。

在修改這個代碼bug後重新生成onnx文件，使用opencv讀取onnx文件做前向計算就不再出現異常錯誤了。

通過以上幾個程序實驗，可以總結出opencv讀取onnx文件做深度學習前向計算的2個坑：

(1) .opencv里的深度學習模塊不支持3維池化計算，解決辦法是修改原始網絡結構，把3維池化轉換成兩個2維池化，重新生成onnx文件

(2) .當神經網絡里有torch.mean和torch.sum這種把4維張量收縮到一個數值的運算時，opencv執行forward會出錯，這時的解決辦法是修改原始網絡結構，在torch.mean的後面加上.item()

在解決這些坑之後，編寫了一套使用opencv做車牌檢測與識別的程序，包含C++和python兩個版本的代碼。使用opencv的dnn模塊做前向計算，後處理模塊是自己使用C++和Python獨立編寫的。

代碼已發布在github上，地址是：https://github.com/hpc203/license-plate-detect-recoginition-opencv

起初在github上看到一個使用DBNet檢測條形碼的程序，不過它是基於pytorch框架做的。於是我編寫一套程序把pytorch模型轉換到onnx文件，使用opencv讀取onnx文件做前向計算。編寫完程序後在運行時沒有出錯，但是最後輸出的結果跟調用pytorch 的輸出結果不一致，並且從可視化結果看，沒有檢測出圖片中的條形碼。這時在看到網上有很多使用onnxruntime部署onnx模型的文章，於是決定使用onnxruntime部署，編寫完程序後運行，選取幾張快遞單圖片測試，結果如下圖所示DBNet檢測到的4個點,圖中綠色的點，紅色的線是把4個連接起來的直線。

並且我還編寫了一個函數比較opencv和onnxruntime的輸出結果，程序代碼和運行結果如下，可以看到在相同輸入，讀取同一個onnx文件的前提下，opencv和onnxruntime的輸出結果竟然不相同。

ONNXRuntime是微軟推出的一款推理框架，用戶可以非常便利的用其運行一個onnx模型。從這個實驗，可以看出相比於opencv庫，onnxruntime庫對onnx模型支持的更好。

我把這套使用DBNet檢測條形碼的程序發布在github上，地址是：https://github.com/hpc203/dbnet-barcode

在第1節講到opencv不支持3維池化，那麼onnxruntime是否支持呢？接着編寫了一個程序探索onnxruntime對3維池化的支持情況，代碼和運行結果如下，可以看到程序報錯了。

查看nn.MaxPool3d的說明文檔，截圖如下，可以看到它的輸入和輸出是5維張量，於是修改上面的代碼，把輸入調整到5維張量。

代碼和運行結果如下，可以看到這時候onnxruntime庫能正常讀取onnx文件，並且它的輸出結果跟pytorch的輸出結果相等。

繼續實驗，把三維池化改作三維卷積，代碼和運行結果如下，可以看到平均差異在小數點後11位，可以忽略不計。

在第1節講到過opencv不支持3維池化，那時候的輸入張量是4維的，如果把輸入張量改成5維的，那麼opencv是否就能進行3維池化計算呢？為此，編寫代碼，驗證這個想法。代碼和運行結果如下，可以看到在cv2.dnn.blobFromImage這行代碼出錯了。

查看cv2.dnn.blobFromImage這個函數的說明文檔，截圖如下，可以看到它的輸入image是4維的，這說明它不支持5維的輸入。

經過這一系列的程序實驗論證，可以看出onnxruntime庫對onnx模型支持的更好。如果深度學習模型有3維池化或3維卷積層，那麼在轉換到onnx文件後，使用onnxruntime部署深度學習是一個不錯的選擇。

不過我在做pytorch導出onnx文件時，還發現了一個問題。在torch.export函數裡有一個輸入參數dynamic_axes，它表示動態的軸，即可變的維度。假如一個神經網絡輸入是動態分辨率的，那麼需要定義dynamic_axes = {'input': {2: 'height', 3: 'width'}, 'output': {2: 'height', 3: 'width'}}，接下來我編寫一個程序來驗證，代碼和運行結果的截圖如下

可以看到，在生成onnx文件後，使用onnxruntime庫讀取，對輸入blob的高增加10個像素單位，在run這一步出錯了。使用opencv讀取onnx文件，代碼和運行結果的截圖如下，可以看到依然出錯了。

通過這個程序實驗，讓人懷疑torch.export函數的輸入參數dynamic_axes是否真的支持動態分辨率輸入的。

以上這些程序實驗是我在編寫算法應用程序時記錄下的一些bug和解決方案的，希望能幫助到深度學習算法開發應用人員少走彎路。

此外，DBNet的官方代碼里提供了轉換到onnx模型文件，於是我依然編寫了一套使用opencv部署DBNet文字檢測的程序，依然是包含C++和Python兩個版本的代碼。官方代碼的模型是在ICDAR場景文本檢測數據集上訓練的，考慮到車牌里也含有文字，我把文章開頭展示的汽車圖片作為輸入，程序檢測結果如下，可以看到依然能檢測到車牌的4個角點，只是不夠準確。如果想要獲得準確的角點定位，可以在車牌數據集上訓練DBNet。

我把使用opencv部署DBNet文字檢測的程序發布在github上，程序依然是包含c++和python兩種版本的實現，地址是：https://github.com/hpc203/dbnet-opencv-cpp-python

掃描下方二維碼關注【集智書童】公眾號，獲取更多實踐項目源碼和論文解讀，非常期待你我的相遇，讓我們以夢為馬，砥礪前行！