部署深度学习模型

载入模型

准备模型

model = models.resnet18(pretrained=True)
model = model.eval().to(device)

获取一组输入

x = torch.randn(1, 3, 256, 256).to(device)

Pytorch模型转ONNX模型

with torch.no_grad():
    torch.onnx.export(
        model,                  # 要转换的模型
        x,                      # 模型的任意一组输入
        'resnet18.onnx',        # 导出的 ONNX 文件名
        opset_version=11,       # ONNX 算子集版本
        input_names=['input'],  # 输入 Tensor 的名称（自己起名字）
        output_names=['output'] # 输出 Tensor 的名称（自己起名字）
    )

验证onnx模型导出成功

import onnx

# 读取 ONNX 模型
onnx_model = onnx.load('resnet18.onnx')

# 检查模型格式是否正确
onnx.checker.check_model(onnx_model)

print('无报错，onnx模型载入成功')

可以使用Netron对onnx模型可视化

应用场景

只需把onnx模型文件发到部署硬件上，并安装 ONNX Runtime 环境，用几行代码就可以运行模型了。

import onnxruntime
import numpy as np
import torch
#载入 onnx 模型，获取 ONNX Runtime 推理器
ort_session = onnxruntime.InferenceSession('resnet18.onnx')

#构造输入，获取输出结果
x = torch.randn(1, 3, 256, 256).numpy()

获取结果

# onnx runtime 输入
ort_inputs = {'input': x}

# onnx runtime 输出
ort_output = ort_session.run(['output'], ort_inputs)[0]

注意，输入输出张量的名称需要和 torch.onnx.export 中设置的输入输出名对应

预处理

input_img = test_transform(img_pil)
input_tensor = input_img.unsqueeze(0).numpy()

ONNX Runtime预测

# ONNX Runtime 输入
ort_inputs = {'input': input_tensor}
# ONNX Runtime 输出
pred_logits = ort_session.run(['output'], ort_inputs)[0]
pred_logits = torch.tensor(pred_logits)

import torch.nn.functional as F
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算