編譯 ONNX 模型

本篇文章譯自英文文檔?Compile ONNX Models — tvm 0.13.dev0 documentation
作者是 Joshua Z. Zhang
更多 TVM 中文文檔可訪問 →TVM 中文站。
本文將介紹如何用 Relay 部署 ONNX 模型。
首先安裝 ONNX 包,最便捷的方法推薦安裝 protobuf 編譯器:
或參考官方網(wǎng)站:
https://github.com/onnx/onnx
加載預訓練的 ONNX 模型
下面示例中的超分辨率模型與 ONNX 教程 中的模型完全相同,跳過 PyTorch 模型的構(gòu)建部分,下載保存的 ONNX 模型:
加載測試圖像
該模型接收大小為 224x224 的單個圖像作為輸入,輸出沿每個軸放大 3 倍的圖像(即大小為 672x672)。為適配輸入的 shape,重新縮放貓圖像,并轉(zhuǎn)換為?YCbCr。然后超分辨率模型應(yīng)用于亮度(Y)通道。
使用 Relay 編譯模型
通常 ONNX 模型將輸入值與參數(shù)值混合在一起,輸入名稱為?1,具體要查看模型文檔來確定完整的輸入和參數(shù)名稱空間。
將 shape 字典傳給?relay.frontend.from_onnx?方法,以便 Relay 知道哪些 ONNX 參數(shù)是輸入,哪些是參數(shù),并提供輸入尺寸的靜態(tài)定義:
輸出結(jié)果:
在 TVM 上執(zhí)行
查看結(jié)果
將輸入和輸出圖像放在一起比對。亮度通道?Y是模型的輸出。將色度通道?Cb?和?Cr?調(diào)整到匹配簡單的雙三次算法,然后將圖像重新組合,并轉(zhuǎn)換回?RGB。

輸出結(jié)果:
注意
ONNX 導入器在導入時默認根據(jù)動態(tài) shape 定義模型,編譯器在編譯時將模型轉(zhuǎn)換為靜態(tài) shape。如果失敗,模型中可能仍存在動態(tài)操作。目前并非所有 TVM 內(nèi)核都支持動態(tài) shape,如果遇到動態(tài)內(nèi)核錯誤,請在 discuss.tvm.apache.org 上提交 issue。
這個特定的模型是用舊版本的 ONNX 構(gòu)建的。在導入階段,ONNX 導入器運行 ONNX 驗證程序(可能拋出屬性類型不匹配的警告)。由于 TVM 支持許多不同的 ONNX 版本,所以 Relay 模型仍然有效。
下載 Python 源代碼:from_onnx.py
下載 Jupyter Notebook:from_onnx.ipynb