onnxruntime模型部署流程

一、將訓練好的模型轉換格式為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使用方法
載入模型：

session = onnxruntime.InferenceSession("./dmnet.onnx")

載入圖片：

img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
tensor = transforms.ToTensor()(img)
tensor = tensor.unsqueeze_(0)

執行推理：
注意：這裡的"input"是和轉onnx格式時的名字對應的。

result = session.run([], {"input": tensor.cpu().numpy()})