图像分类----模型部署
代码流程
把原生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)