【模型部署入门一】：Pytorch图像分类模型转换ImageNet1000类预训练模型转ONNX

入门学习主要是跟随同济子豪兄，很感谢该博主，本博客代码主体是子豪兄的，我只是总结加工整理记录。

子豪兄对应学习视频链接地址为：ImageNet1000类预训练模型转ONNX

我自己总结的教程中代码使用任何python编译环境都可以直接打开使用，子豪兄视屏教程中的源码需要使用Jupyter打开，关于Jupyter的安装及使用，推荐参考博客：Jupyter Notebook的安装及使用教程

一、安装配置环境

如果还没有安装CUDA，CUDNN，推荐安装教程：安装CUDA+CUDNN教程

1.1 安装Pytorch

pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html

1.2 安装ONNX

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

1.3 安装推理引擎ONNX Runtime

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

1.4 安装其它包

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

1.5 检验是否安装成功

import torch
print('PyTorch 版本', torch.__version__)

import onnx
print('ONNX 版本', onnx.__version__)

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

1.6 检验CUDA是否被调用

import torch

# 检查GPU是否可用
if torch.cuda.is_available():
    print("PyTorch is using GPU")
else:
    print("PyTorch is using CPU")

# 检查当前使用的GPU设备的索引
try:
    device_index = torch.cuda.current_device()
    print("Current GPU index: ", device_index)
except RuntimeError:
    print("No GPU available, using CPU")

# 检查当前系统中可用的GPU设备数量
num_devices = torch.cuda.device_count()
print("Number of available GPUs: ", num_devices)

如果电脑GPU能被正常调用，运行代码后的输出如下：

在这里插入图片描述

二、ONNX官网

ONNX官网地址：ONNX

打开后的样子如下:

在这里插入图片描述

查看算子版本：

在这里插入图片描述

三、ImageNet1000类预训练模型转ONNX

下面代码中有详细注释，学者可直接运行使用。

import torch
from torchvision import models
import onnx

## 有GPU就用GPU，无就用CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
print(torch.cuda.is_available())
print(device)

## 载入ImageNet预训练Pytorch图像分类模型
model = models.resnet18(pretrained=True)      # 从PyTorch的模型库中加载一个预训练的ResNet-18模型。pretrained=True表示加载的模型已经在ImageNet数据集上进行了预训练
model = model.eval().to(device)               # 调用eval()方法将模型设置为评估模式。在评估模式下，模型的某些层（如Dropout层和BatchNorm层）会改变行为。然后，to(device)方法将模型移动到指定的设备上

## 构造一个输入图像Tensor
x = torch.randn(1,3,256,256).to(device)
print(x.shape)     # torch.Size([1, 3, 256, 256])

## 输入Pytorch模型推理预测，获得1000个类别的预测结果
output = model(x)  # 将随机张量x输入到模型中运行前向推理  output就是1000个类别的预测分数
print(output)
print(output.shape)     # torch.Size([1, 1000])

## Pytorch模型转ONNX格式
with torch.no_grad():
    torch.onnx.export(
        model,                      # 要转换的模型
        x,                          # 模型的任意一组输入 张量的维度一定要和模型指定输入的维度一致
        'resnet18_imagenet.onnx',  # 导出的ONNX文件名，自定义名称
        opset_version = 11,         # ONNX 算子集版本，ONNC算子官网地址：https://onnx.ai/
        input_names = ['input'],    # 输入Tensor的名称   # 自定义名称
        output_names = ['output']   # 输出Tensor的名称   # 自定义名称
    )

## 验证导出的onnx模型是否导出成功
onnx_model = onnx.load('resnet18_imagenet.onnx')   # 读取上一步导出的ONNX模型
onnx.checker.check_model(onnx_model)               # 检测模型格式是否正确

print("恭喜你，ONNX模型载入成功！")

## 以可读形式打印计算图
print(onnx.helper.printable_graph(onnx_model.graph))

## 网页端打开
# 网址：https://netron.app/