Pytorch图像分类模型模型实时在线验证代码

1.训练并保存自己的模型

保存的模型格式为：XXX.pth

torch.save(model, "./weight/last.pth")if best_acc <(validation_acc / len_val):torch.save(model, "./weight/best.pth")

2.转化为ONNX格式

2.1环境安装（window10）

pip install onnx
pip install onnxruntime#验证安装配置是否成功
import torch
print('PyTorch 版本', torch.__version__)import onnx
print('ONNX 版本', onnx.__version__)import onnxruntime as ort
print('ONNX Runtime 版本', ort.__version__)

2.2.pth格式转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)model = torch.load('best.pth')
model = model.eval().to(device)
x = torch.randn(1, 3, 256, 256).to(device)  #这里要构造一个数据，保证和自己输入的图片大小一致3*256*256
output = model(x)  #output.shape = torch.Size([1, 10])  这是一个10分类问题#Pytorch模型转ONNX模型
x = torch.randn(1, 3, 256, 256).to(device)with torch.no_grad():torch.onnx.export(model,                   # 要转换的模型x,                       # 模型的任意一组输入'best.onnx', # 导出的 ONNX 文件名opset_version=11,        # ONNX 算子集版本input_names=['input'],   # 输入 Tensor 的名称（自己起名字）output_names=['output']  # 输出 Tensor 的名称（自己起名字）import onnx# 读取 ONNX 模型
onnx_model = onnx.load('resnet18_fruit30.onnx')# 检查模型格式是否正确
onnx.checker.check_model(onnx_model)
print('无报错，onnx模型载入成功')

这是project中就出现了“best.onnx”文件，表示转化ONNX格式成功！

3.可视化实时检测

3.1在PC电脑端查看

3.1.1环境安装（待补充）

pip install onnxruntime
需要提前保存一个类别ID和类别名称对应的文件

3.1.2 摄像头实时捕捉并分类

import onnxruntime
import torch
from torchvision import transforms
import torch.nn.functional as F
import pandas as pd
import numpy as np
from PIL import Image, ImageFont, ImageDraw
import matplotlib.pyplot as plt# 导入中文字体，指定字号
font = ImageFont.truetype('SimHei.ttf', 32)#载入ONNX模型，获取ONNX Runtime推力器
ort_session = onnxruntime.InferenceSession('best.onnx')#载入类别和ID对应字典
idx_to_labels = np.load('idx_to_labels.npy', allow_pickle=True).item()# 测试集图像预处理-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])])