一、将訓練好的模型轉換格式為ONNX格式
例如pytorch模型轉換:
def torch2onnx(model, save_path):
"""
:param model:
:param save_path: XXX/XXX.onnx
:return:
"""
model.eval()
data = torch.rand(1, 3, 224, 224)
input_names = ["input"]
output_names = ["out"]
torch.onnx._export(model, data, save_path, export_params=True, opset_version=11, input_names=input_names, output_names=output_names)
print("torch2onnx finish.")
支援動态形狀的輸入和輸出:
def torch2onnx_dynamic(model, save_path):
"""
:param model:
:param save_path: XXX/XXX.onnx
:return:
"""
model.eval()
data = torch.rand(1, 3, 224, 224)
input_names = ["input"] # ncnn需要
output_names = ["out"] # ncnn需要
torch.onnx._export(model, data, save_path, export_params=True, opset_version=11, input_names=input_names,
output_names=output_names, dynamic_axes={'input': [2, 3], 'out': [2, 3]})
print("torch2onnx finish.")
二、安裝onnxruntime
注意:onnxruntime-gpu版本在0.4以上時需要CUDA 10
pip install onnxruntime
pip install onnxruntime-gpu
onnxruntime幫助文檔:
https://microsoft.github.io/onnxruntime/python/tutorial.html
三、onnxruntime使用方法
加載模型:
加載圖檔:
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
tensor = transforms.ToTensor()(img)
tensor = tensor.unsqueeze_(0)
執行推理:
注意:這裡的"input"是和轉onnx格式時的名字對應的。