Pytorch模型的部署

教程参考B站Up主：同济子豪兄

视频：https://space.bilibili.com/1900783

---------------------------------------------------------------------------------------------------------------------------------

一、安装所需要的环境

#  安装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



###########################################################################################
#  验证Pytorch的版本
import torch
print('Pytorch版本为：',torch.__version__)
#  验证onnx的版本
import onnx
print("ONNX的版本：",onnx.__version__)
#  验证ONNXRuntime的版本
import onnxruntime as ort
print('ONNXRuntime的版本：',ort.__version__)

二、

import torch
from torchvision import models
import onnx

#  确认在推理过程中使用GPU还是CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device:',device)

#  加载预训练模型，  pretrained = True会自动下载该模型
model = models.resnet18(pretrained = True)
model = model.eval().to(device)

#  1表示batch_size
x = torch.randn(1, 3, 256, 256).to(device)
output = model(x)

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


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

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

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