教程参考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模型载入成功')