Pytorch图像分类模型转ONNX（同济子豪兄学习笔记）

安装配置环境
代码运行云GPU平台：公众号 人工智能小技巧 回复 gpu
同济子豪兄 2022-8-22 2023-4-28 2023-5-8
安装 Pytorch

pip3 install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113

安装 ONNX

pip install onnx -i https://pypi.tuna.tsinghua.edu.cn/simple

安装推理引擎 ONNX Runtime

pip install onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple

安装其它第三方工具包

pip install numpy pandas matplotlib tqdm opencv-python pillow -i https://pypi.tuna.tsinghua.edu.cn/simple

验证安装配置成功

import torch
import onnx
import onnxruntime as ort
print('PyTorch 版本', torch.__version__)
print('ONNX 版本', onnx.__version__)
print('ONNX Runtime 版本', ort.__version__)

Pytorch图像分类模型转ONNX-ImageNet1000类
把Pytorch预训练ImageNet图像分类模型，导出为ONNX格式，用于后续在推理引擎上部署。
代码运行云GPU平台：公众号 人工智能小技巧 回复 gpu
同济子豪兄 2022-8-22 2023-4-28 2023-5-8

import torch
from torchvision import models
import onnx

device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device:',device)

# 载入ImageNet预训练PyTorch图像分类模型
model=models.resnet18(pretrained=True)
model=model.eval().to(device)
# 构造一个输入图像Tensor
x=torch.rand(1,3,256,256).to(device)
print(x.shape)
# 输入PyTorch模型推理预测，获得1000个类别预测结果
output=model(x)
print(output.shape)
# PyTorch模型转ONNX格式
with torch.no_grad():
    torch.onnx.export(
        model,                    # 要转换的模型
        x,                        # 模型的任意一组输入
        'resnet18_imagenet.onnx', # 导出的ONNX文件名
        input_names=['input'],    # 输入Tensor的名称（自己起名字）
        output_names=['output'],  # 输出Tensor的名称（自己起名字）
        opset_version=11,         # ONNX算子集版本
    )
# 验证onnx模型导出成功
# 读取ONNX模型
onnx_model=onnx.load('resnet18_imagenet.onnx')
# 检查模型格式是否正确
onnx.checker.check_model(onnx_model)
print('无报错，onnx模型导入成功!')
# 以可读的形式打印计算图
print(onnx.helper.printable_graph(onnx_model.graph))
#使用Netron可视化模型结构
#Netron：https://netron.app
#视频教程：https://www.bilibili.com/video/BV1TV4y1P7AP

PyTorch图像分类模型转ONNX-花分类

import torch
from torchvision import models
import onnx

device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device:',device)

# 导入训练好的模型
model=torch.load('MobileNetV2.pth')
model=model.eval().to(device)
# 构造一个输入图像Tensor
x=torch.rand(1,3,224,224).to(device)
print(x.shape)
# 输入PyTorch模型推理预测，获得N个类别预测结果
output=model(x)
print(output.shape)
# PyTorch模型转ONNX格式
with torch.no_grad():
    torch.onnx.export(
        model,                    # 要转换的模型
        x,                        # 模型的任意一组输入
        'mobilenetv2_flowers5.onnx', # 导出的ONNX文件名
        input_names=['input'],    # 输入Tensor的名称（自己起名字）
        output_names=['output'],  # 输出Tensor的名称（自己起名字）
        opset_version=11,         # ONNX算子集版本
    )
# 验证onnx模型导出成功
# 读取ONNX模型
onnx_model=onnx.load('mobilenetv2_flowers5.onnx')
# 检查模型格式是否正确
onnx.checker.check_model(onnx_model)
print('无报错，onnx模型导入成功!')
# 以可读的形式打印计算图
print(onnx.helper.printable_graph(onnx_model.graph))
# 登录netron.app在线模型可视化网页打开可视化模型结构

扩展阅读
【同济子豪兄】两天搞定图像分类、目标检测、语义分割、实例分割、关键点检测毕业设计：https://github.com/TommyZihao/Train_Custom_Dataset
ONNX主页： https://onnx.ai
ONNX Github主页： https://github.com/onnx/onnx
ONNX支持的深度学习框架：https://onnx.ai/supported-tools.html
ONNX支持的计算平台：https://onnx.ai/supported-tools.html
ONNX算子：https://onnx.ai/onnx/operators
模型在线可视化网站 Netron：https://netron.app
Netron视频教程：https://www.bilibili.com/video/BV1TV4y1P7AP
周奕帆博客
模型部署入门教程（一）：模型部署简介 https://zhuanlan.zhihu.com/p/477743341
模型部署入门教程（二）：解决模型部署中的难题 https://zhuanlan.zhihu.com/p/479290520
模型部署入门教程（三）：PyTorch 转 ONNX 详解 https://zhuanlan.zhihu.com/p/498425043