onnx模型轉化與跨框架部署詳解+程式碼【pytorch版】
文章目錄
引言
目前訓練階段最流行的框架是pytorch和tensorflow，訓練完的模型常常需要整合到現有的應用程式中、部署到不同型別的平臺上（比如雲端、邊端等）。如果要在每個平臺上實現所有模型的框架，會極大增加環境的複雜性，優化不同框架和硬體的所有組合非常耗時。
所以需要一種通用的 解決方案，來整合、部署、優化不同框架訓練出的模型。ONNX就是為了解決這個問題而誕生。
現在常用的跨框架部署方案有：
- pytorch模型（訓練框架）→onnx模型→onnxruntime部署（推理引擎）
- pytorch模型（訓練框架）→onnx模型→TensorRT部署（推理引擎）
- pytorch模型（訓練框架）→MMdnn模型轉換框架→tensorflow/caffe/MXNet/CoreML部署（推理框架）
（其他方案歡迎評論區補充~）
下一篇文章具體介紹下TensorRT部署onnx模型的兩種具體方案。
基礎概念
onnx：跨框架的模型表達標準
開放式神經網路交換（ONNX），為深度學習和傳統ML的AI模型提供了一種開源格式。它定義了一個可擴充套件的計算圖形模型，以及內建運算子和標準資料型別的定義。
當前，ONNX專注於推理所需的功能（暫不支援訓練的相關功能）。
Microsoft 和合作夥伴社群建立了 ONNX 作為表示機器學習模型的開放標準。 許多框架（包括 TensorFlow、PyTorch、SciKit-Learn、Keras、Chainer、MXNet、MATLAB 和 SparkML）中的模型都可以匯出或轉換為標準 ONNX 格式。 模型採用 ONNX 格式後，可在各種平臺和裝置上執行。
onnxruntime：部署模型的推理引擎
onnxruntime是一種用於將 ONNX 模型部署到生產環境的高效能推理引擎，與許多流行的ML / DNN框架相容，包括PyTorch，TensorFlow / Keras，scikit-learn等。
-
針對雲和 Edge 進行了優化，適用於 Linux、Windows 和 Mac。
-
使用 C++ 編寫，還包含 C、Python、C#、Java 和 Javascript (Node.js) API，可在各種環境中使用。
-
同時支援 DNN 和傳統 ML 模型，並與不同硬體上的加速器（例如，NVidia GPU 上的 TensorRT、Intel 處理器上的 OpenVINO、Windows 上的 DirectML 等）整合。
採用ONNX Runtime的主要優點是：
- 提高各種ML模型的推理效能
- 減少訓練大型模型的時間和成本
- 使用Python進行訓練，但可以部署到C＃/ C ++ / Java應用程式中
- 在不同的硬體和作業系統上執行
- 支援在多個不同框架中建立的模型
API文件：https://www.onnxruntime.ai/python/index.html
示例程式碼
0）安裝onnx和onnxruntime
# pytorch一般自帶onnx
conda install -c conda-forge onnx
# onnxruntime的其他版本有可能報錯ImportError: cannot import name 'get_all_providers'
pip install onnxruntime # CPU build
pip install onnxruntime-gpu # GPU build
如果需要編譯安裝，可以參照ONNX Runtime 的GitHub倉庫地址 上的說明。
需要注意的是：
- 預設安裝方式只支援Python3
- windows下編譯安裝需要Visual C++ 2019 runtime，且只支援win10及以上版本
1）pytorch模型轉onnx模型
利用torch.onnx.export匯出模型到ONNX格式。
import torch.onnx
import torchvision
# Standard ImageNet input - 3 channels, 224x224,
# values don't matter as we care about network structure.
# But they can also be real inputs.
dummy_input = torch.randn(1, 3, 224, 224)
# Obtain your model, it can be also constructed in your script explicitly
model = torchvision.models.alexnet(pretrained=True)
# Invoke export
torch.onnx.export(model, dummy_input, "alexnet.onnx")
2）onnx模型檢驗
檢查onnx模型，並列印模型的結構表示。onnx模型還可以通過視覺化工具（如 Netron）檢視。
import onnx
# Load the ONNX model
model = onnx.load("alexnet.onnx")
# Check that the IR is well formed
onnx.checker.check_model(model)
# Print a human readable representation of the graph
print(onnx.helper.printable_graph(model.graph))
3）呼叫ONNX Runtime測試輸入圖片
採用python API，載入一張圖片測試我們轉化後的onnx模型。
import numpy as np # we're going to use numpy to process input and output data
import onnxruntime # to inference ONNX models, we use the ONNX Runtime
import time
from PIL import Image
def load_labels(path):
with open(path) as f:
data = json.load(f)
return np.asarray(data)
# 影象預處理
def preprocess(input_data):
# convert the input data into the float32 input
img_data = input_data.astype('float32')
#normalize
mean_vec = np.array([0.485, 0.456, 0.406])
stddev_vec = np.array([0.229, 0.224, 0.225])
norm_img_data = np.zeros(img_data.shape).astype('float32')
for i in range(img_data.shape[0]):
norm_img_data[i,:,:] = (img_data[i,:,:]/255 - mean_vec[i]) / stddev_vec[i]
#add batch channel
norm_img_data = norm_img_data.reshape(1, 3, 224, 224).astype('float32')
return norm_img_data
def softmax(x):
x = x.reshape(-1)
e_x = np.exp(x - np.max(x))
return e_x / e_x.sum(axis=0)
def postprocess(result):
return softmax(np.array(result)).tolist()
# Load the raw image
img = Image.open("D://Pic//cat.jpg")
img = img.resize((224, 224), Image.BILINEAR)
print("Image size: ", img.size)
image_data = np.array(img).transpose(2, 0, 1)
input_data = preprocess(image_data)
# Run the model on the backend
session = onnxruntime.InferenceSession('.//alexnet.onnx', None)
# get the name of the first input of the model
input_name = session.get_inputs()[0].name
print('Input Name:', input_name)
# Inference
start = time.time()
raw_result = session.run([], {input_name: input_data})
end = time.time()
res = postprocess(raw_result)
inference_time = np.round((end - start) * 1000, 2)
idx = np.argmax(res)
print('========================================')
print('Final top prediction is: %d'% idx)
print('========================================')
print('========================================')
print('Inference time: ' + str(inference_time) + " ms")
print('========================================')
參考教程
[1] https://github.com/onnx/tutorials/blob/master/tutorials/PytorchOnnxExport.ipynb
[2] https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html
[3] https://github.com/onnx/onnx-docker/blob/master/onnx-ecosystem/inference_demos/resnet50_modelzoo_onnxruntime_inference.ipynb