参考文章：
【netron】模型可视化工具netron-CSDN博客
Pytorch图像分类模型部署-ONNX Runtime本地终端推理_哔哩哔哩_bilibili
使用netron可视化模型结构
1）使用在线版
浏览器访问：Netron
点击 “Open Model” 按钮，选择要可视化的模型文件即可
2）下载本地版
终端进行安装： pip install netron
安装完成后，在脚本中 调用包 import netron
运行程序 netron.start("model.onnx")， 会自动打开浏览器进行可视化 （最后有例子）
我习惯用 pytorch，但是 netron 对 pytorch 的 .pt 和 .pth 文件不是很友好，所以，我都是先转换为 onnx 格式，再进行可视化，下面举例。
简化可视化模型：
ONNX模型转换及使用指南-CSDN博客
onnx的概念
训练的归训练，部署的归部署。 onnx的存在极大地降低了部署的难度。
安装配置环境：
安装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
print('PyTorch 版本', torch.__version__)import onnx
print('ONNX 版本', onnx.__version__)import onnxruntime as ort
print('ONNX Runtime 版本', ort.__version__)转ONNX模型文件
导入工具包
import torch
from torchvision import models
# 有 GPU 就用 GPU，没有就用 CPU
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.randn(1, 3, 256, 256).to(device)输入Pytorch模型推理预测，获得1000个类别的预测结果
output = model(x)output.shapePytorch模型转ONNX格式
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 的名称（自己起名字）
) 我只修改了自己代码的这个地方↓
# '''pth模型文件转为onnx格式'''
with torch.no_grad():
torch.onnx.export(
model, # 要转换的模型
img, # 模型的任意一组输入
'resnet34_classification.onnx', # 导出的 ONNX 文件名
opset_version=11, # ONNX 算子集版本
input_names=['input'], # 输入 Tensor 的名称（自己起名字）
output_names=['output'] # 输出 Tensor 的名称（自己起名字）
) 验证onnx模型导出成功
import onnx
# 读取 ONNX 模型
onnx_model = onnx.load('resnet18_imagenet.onnx')
# 检查模型格式是否正确
onnx.checker.check_model(onnx_model)
print('无报错，onnx模型载入成功')推理引擎ONNX Runtime部署-预测单张图像
导入工具包
import onnxruntime
import numpy as np
import torch
import torch.nn.functional as F
import pandas as pd载入 onnx 模型，获取 ONNX Runtime 推理器
ort_session = onnxruntime.InferenceSession('resnet18_imagenet.onnx')构造随机输入，获取输出结果
x = torch.randn(1, 3, 256, 256).numpy()x.shape # (1, 3, 256, 256)# onnx runtime 输入
ort_inputs = {'input': x}
# onnx runtime 输出
ort_output = ort_session.run(['output'], ort_inputs)[0]注意，输入输出张量的名称需要和 torch.onnx.export 中设置的输入输出名对应
ort_output.shape # (1, 1000)载入一张真正的测试图像
img_path = 'banana1.jpg'# 用 pillow 载入
from PIL import Image
img_pil = Image.open(img_path)预处理函数
from torchvision import transforms
# 测试集图像预处理-RCTN：缩放裁剪、转 Tensor、归一化
test_transform = transforms.Compose([transforms.Resize(256),
transforms.CenterCrop(256),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])运行预处理
input_img = test_transform(img_pil)input_img.shape # torch.Size([3, 256, 256])input_tensor = input_img.unsqueeze(0).numpy()input_tensor.shape # (1, 3, 256, 256)推理预测
# ONNX Runtime 输入
ort_inputs = {'input': input_tensor}
# ONNX Runtime 输出
pred_logits = ort_session.run(['output'], ort_inputs)[0]
pred_logits = torch.tensor(pred_logits)pred_logits.shape # torch.Size([1, 1000]) 因为ImageNet有1000类# 对 logit 分数做 softmax 运算，得到置信度概率
pred_softmax = F.softmax(pred_logits, dim=1) pred_softmax.shape # torch.Size([1, 1000])解析预测结果
# 取置信度最高的前 n 个结果
n = 3
top_n = torch.topk(pred_softmax, n)
top_n
'''torch.return_types.topk(
values=tensor([[9.9669e-01, 2.6005e-03, 3.0254e-04]]),
indices=tensor([[954, 939, 941]]))'''
# 预测类别
pred_ids = top_n.indices.numpy()[0]
pred_ids
'''array([954, 939, 941])'''用自己转换好的ONNX模型文件类型预测
# 加载ONNX模型
onnx_model_path = 'resnet34_classification.onnx'
onnx_model = onnx.load(onnx_model_path)
# 创建ONNX Runtime会话
ort_session = onnxruntime.InferenceSession(onnx_model_path)
# 将PyTorch Tensor转换为NumPy数组
input_data = img.numpy()
# 运行预测
ort_inputs = {ort_session.get_inputs()[0].name: input_data}
ort_outputs = ort_session.run(['output'], ort_inputs)
# 获取预测结果
output = ort_outputs[0]
predict = torch.from_numpy(output)
predict_cla = torch.argmax(predict).item()
grading_result = int(class_indict[str(predict_cla)])
end_time = time.time()
# 打印预测时间
prediction_time = end_time - start_time
print("模型预测时间：{:.4f}秒".format(prediction_time))
return grading_result注意：相比于直接用pt文件预测图片，输入的图像转成numpy的数据格式
# 将PyTorch Tensor转换为NumPy数组
input_data = img.numpy()
