图像分类----模型部署

图像分类----模型部署

代码流程

       把原生Pytorch训练得到的图像分类模型，导出为ONNX格式，用于后续在ONNX Runtime推理引擎上部署。

1.引入库

import onnx
import torch
from torchvision import models
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

2.载入预训练分类模型

model = models.resnet18(pretrained=True)
model = model.eval().to(device)
x = torch.randn(1, 3, 256, 256).to(device)
output = model(x)

3.Pytorch模型转ONNX模型

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

4.验证onnx模型导出成功

onnx_model = onnx.load('resnet18.onnx')
onnx.checker.check_model(onnx_model)