實踐教程 | 輕松入門模型轉換和可視化
共 5993字,需瀏覽 12分鐘
·
2024-07-15 22:00
來源 | CV研習社
編輯 | 極市平臺
極市導讀
本文給大家介紹一個模型轉換格式ONNX和可視化工具Netron。ONNX是微軟設計的一種多平臺的通用文件格式,幫助算法人員進行模型部署和框架之間相互轉換。而Netron是一款老牌的輕量化模型可視化工具,支持多種開源框架,使用上比TensorBoard簡單直接。 >>加入極市CV技術交流群,走在計算機視覺的最前沿
1 不同框架模型需要轉換嗎?
目前開源的深度學習框架有很多,當我們在Github上搜索一個通用網絡時,會發(fā)現各種各樣的版本實現。但是深度學習方法的本質都是分為訓練和部署兩個環(huán)節(jié),如果以目標為導向,我們關注算法部署在哪兒,用什么部署?是x86上采樣TensorRT部署還是ARM上采用端框架NCNN運行,又或者采樣某個硬件廠商提供的SDK做推理。
假設童鞋A有一個檢測任務需要部署在TX2的開發(fā)板上,基于TensorFlow框架把yolov5改改,訓練得到了一份pb權重,接下來將pb轉換成TensorRT使用的uff文件做推理。
如果A將框架換成Pytorch,按照上述流程得到了一份pt的權重,然后調用TensorRT的SDK轉換成uff格式,看上去好像也很順暢沒什么問題,它們的區(qū)別就在于訓練生成的權重文件格式不同,不過反正最終都轉換成uff格式。
實際操作中或者部署在其他移動端時,我們會發(fā)現瓶頸在于如何才能更友好的將PC端訓練的權重文件轉換成部署端需要的格式。Nividia的開發(fā)板雖然基于CUDA加速,生態(tài)鏈完整,但是對不同的神經網絡框架支持的算子力度也有區(qū)別,如果訓練框架是Pytorch,在TensorRT上部署,如下圖所示理論上有三條路(這里Caffe的存在是因為TRT對它的支持更友好,可以根據需求刪除):
直接調用TRT提供的SDK轉換接口,很大可能遇到不支持的操作算子,在Nividia顯卡上還能通過CUDA編程自定義算法,換成其他硬件可能直接就不支持自定義功能。
既然直接轉換走不通,那么能否將Pytorch先轉成Caffe,再由Caffe轉到TRT呢?
這條路也存在問題,不同神經網絡框架在一些定義上存在差異,比如Pytorch的反卷積自帶兩次padding;最大池化時Pytorch向下取整,而Caffe向上取整等等。細節(jié)上的問題會導致轉換后輸出不同。
所以我們可以看到圖上那條更迂回的路,經過ONNX這個較為通用的格式,因為Pytorch—>ONNX—>Caffe—>TensorRT從部署角度看才更可行。
2 什么是ONNX?
ONNX是一套神經網絡模型的開放格式,支持Core ML, PaddlePaddle, SNPE, MXNet, Caffe2, PyTorch, TensorFlow, CNTK 等框架。
ONNX的安裝比較簡單,在Ubuntu上直接pip install onnx即可,從圖上可以看出它是依賴Protobuf庫的,因為它內部采用的是Protobuf作為其數據存儲和傳輸。
ONNX是一種可擴展的計算圖模型、一系列內置的運算單元和標準數據類型。每一個計算流圖都定義為由節(jié)點組成的列表,并構建有向無環(huán)圖。其中每一個節(jié)點都有一個或多個輸入與輸出,每一個節(jié)點稱之為一個 運算單元。這相當于一種通用的計算圖,不同深度學習框架構建的計算圖都能轉化為它。
為什么要把訓練的模型轉換成ONNX?
ONNX的作用是在五花八門的訓練框架和屈指可數的部署方案之間建立橋梁。它是一個中間環(huán)節(jié),方便模型的部署和在各個框架之間轉移。
試下如果你電腦上是Tensorflow的環(huán)境,卻發(fā)現最新的開源模型是用Pytorch寫的,你想快速測試效果怎么辦?
-
采用TensorFlow重寫一遍嗎?肯定不行,先不說實現后性能是否一致,單就花費的時間成本就很高。
-
那么配置一份Pytorch的運行環(huán)境跑一下結果?看起來可行,但是下一次是Caffe,CoreML呢?難不成都配置一遍嗎!
-
再想想辦法,寫個腳本將Pytorch的權重轉成TensorFlow的pb文件?好像也可以,但是實際操作下來會發(fā)現不同框架直接的直接轉換有許多細節(jié)上的不統一。
最后大家商量一個共同使用的中間接口過度不就可以了,微軟開發(fā)的ONNX因此出現了。
還有一種部署的場景是當用某個開源框架訓練的模型文件轉換到ONNX后,還需要將它再次轉換,比如轉換到TensorRT后做前向傳播。為什么要轉換這么多次呢?
我們反過來看就清楚了,目標平臺支持哪個框架力度大,我們就通過ONNX做為中轉站將當前的框架模型轉換到支持算子多的框架,為了讓算法在移動端部署成功。
如何轉換ONNX?
這里以tensorflow為例,從tf轉到onnx還處于實驗階段,我們可以依靠外部工具進行轉換。
-
首先第一步需要安裝tf2onnx包,即pip install -U tf2onnx。
-
然后第二步將ckpt文件凍結成pb文件,這里是因為ckpt和meta文件分別權重信息和拓撲結構分開存儲,而且固化到pb中是將變量凍結成常量形式。TensorFlow官方已經提過了相關工具freeze_graph。
-
最后第三部執(zhí)行tf2onnx.convert命令,即:
python3 -m tf2onnx.convert --input model.pb --inputs input_name:0[1,640,640,3] --outputs output_name:0 --output model.onnx
其中input_name是網絡輸入的名稱,output_name是網絡輸出的名稱。
3 網絡可視化工具
隨著深度學習從2016年的爆發(fā)到現在已有5年時間,市面上深度學習框架也是百花齊放,有TensorFlow、Caffe、Pytorch、NCNN、CoreML、MXNet、ArmNN等等,再加上各大硬件廠商集成的SDK(SNPE、TIDL、TensorRT等),加起來不少于小幾十種。
但是好用的模型可視化工具卻相對非常較少,其中TensorBoard應該算煉丹師們較為熟悉的一種工具,它能夠查看網絡結構,損失的變化,準確率指標等變化情況。功能比較全面,界面如下圖所示:
TensorBoard安裝和使用
TensorBoard是隨著TensorFlow一起安裝的,這里我們不再啰嗦。使用時分為五步:
-
指定文件用于保存Tensorboard的圖:writer = tf.summary.FileWriter(log_dir=log_dir)
-
增加需要觀測的變量:
tf.summary.histogram('weights',weight)
tf.summary.scalar('loss',loss) ......
-
合并所有待觀察的數據:
merged = tf.summary.merge_all()
-
運行命令:tensorboard --logdir=./logs
-
在瀏覽器輸入網址:http://localhost:6006,即可查看生成圖和參數
在第2步中,具體TensorBoard可以觀察哪些東西呢?
-
標量Scalars:存儲和顯示諸如學習率和損失等單個值的變化趨勢 -
圖片Images:對于輸入是圖像的模型,顯示某一步輸入給模型的圖像 -
計算圖Graph:顯示代碼中定義的計算圖,也可以顯示包括每個節(jié)點的計算時間、內存使用等情況 -
數據分布Distribution:顯示模型參數隨迭代次數的變化情況 -
直方圖Histograms:顯示模型參數隨迭代次數的變化情況 -
嵌入向量Embeddings:在3D或者2D圖中展示高維數據
下面我們再推薦另一款輕量級的可視化工具Netron。
Netron安裝和使用
Netron支持兩種方式打開,其一是瀏覽器訪問,無需任何安裝。登錄IP地址:https://netron.app/
點擊Open Model選中你要可視化的模型格式,即可看到如下界面:
圖中可以看到中間區(qū)域是網絡結構的描述,左側是菜單欄,可選擇顯示節(jié)點屬性、名稱等功能,右側是選中的節(jié)點特性,包括類型、名稱、輸入、輸出等。
另一種打開方式需要進行軟件安裝,不同操作系統安裝方式不同,官網上有各個操作系統下的軟件包鏈接:
客戶端模式下加載權重文件,感覺比瀏覽器版本更清晰一些:
Netron支持的框架
Netron目前支持的框架有:
-
Keras的h5文件 -
Caffe的caffemodel和prototxt文件 -
TensorFlow Lite的tflite文件 -
ONNX的onnx文件 -
Core ML的mlmodel文件 -
NCNN的param文件
還有一些處于實驗階段,還不夠穩(wěn)定,如:
-
Pytorch的pt文件 -
TensorFlow的pb、ckpt文件 -
CNTK的model、cntk文件 -
Arm NN的armnn文件 -
Scikit-learn的pkl文件 -
Deeplearning4j的zip文件
從使用角度來說,Netron更適合在拿到一份陌生的網絡時查看其結構;而TensorBoard更適合在訓練的過程中實時關注訓練參數的變化,以便及時作出調整。
公眾號后臺回復“數據集”獲取60+深度學習數據集下載~
# CV技術社群邀請函 #
備注:姓名-學校/公司-研究方向-城市(如:小極-北大-目標檢測-深圳)
即可申請加入極市目標檢測/圖像分割/工業(yè)檢測/人臉/醫(yī)學影像/3D/SLAM/自動駕駛/超分辨率/姿態(tài)估計/ReID/GAN/圖像增強/OCR/視頻理解等技術交流群
每月大咖直播分享、真實項目需求對接、求職內推、算法競賽、干貨資訊匯總、與 10000+來自港科大、北大、清華、中科院、CMU、騰訊、百度等名校名企視覺開發(fā)者互動交流~
