七、图像分类模型的部署（Datawhale组队学习）
文章目录
本文内容为
同济子豪兄图像分类系列视频的学习笔记，
项目参考代码。本文使用ONNX-ONNX Runtime部署我们的模型。
前言
ONNX简介
ONNX是一种针对机器学习所设计的开放式的文件格式，用于存储训练好的模型。它使得不同的人工智能框架（如Pytorch, TensorFlow等）可以采用相同格式存储模型数据并交互。
应用场景
后面的代码在需要部署的硬件上运行，只需把onnx模型文件发到部署硬件上，并安装 ONNX Runtime 环境，用几行代码就可以运行模型了。
pip install onnx onnxruntime
部署ImageNet预训练图像分类模型
导出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预训练图像分类模型
model = models.resnet18(pretrained=True)
model = model.eval().to(device)
Pytorch模型转ONNX模型
x = torch.randn(1, 3, 256, 256).to(device)
with torch.no_grad():
torch.onnx.export(
model, # 要转换的模型
x, # 模型的任意一组输入
'resnet18.onnx', # 导出的 ONNX 文件名
opset_version=11, # ONNX 算子集版本
input_names=['input'], # 输入 Tensor 的名称（自己起名字）
output_names=['output'] # 输出 Tensor 的名称（自己起名字）
)
验证onnx模型导出成功
import onnx
# 读取 ONNX 模型
onnx_model = onnx.load('resnet18.onnx')
# 检查模型格式是否正确
onnx.checker.check_model(onnx_model)
print('无报错，onnx模型载入成功')
#以可读形式打印计算图
#print(onnx.helper.printable_graph(onnx_model.graph))
无报错，onnx模型载入成功
使用Netron对onnx模型可视化
关于如何理解这个网络我们可以参考子豪兄的【精读AI论文】ResNet深度残差网络，一下为ResNet网络的简要介绍。
Resnet将一个模块的输入分为两条路。右边这条路称为短路连接，这个连接将输入原封不动的传入到输出。左边的这条路是两层的神经网络，这两层神经网络不用拟合复杂的底层映射，只用你和原来输入的基础上进行的偏移和修改（即残差）就可以了。最后将残差和恒等映射相加再用Relu激活函数处理。
推理引擎ONNX Runtime部署-预测单张图像
使用推理引擎 ONNX Runtime，读取 onnx 格式的模型文件，对单张图像文件进行预测。
import onnxruntime
import numpy as np
import torch
前期准备
载入 onnx 模型，获取 ONNX Runtime 推理器
ort_session = onnxruntime.InferenceSession('resnet18.onnx')
预处理
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])
])
载入测试图像并进行预处理
img_path = 'banana1.jpg'
# 用 pillow 载入
from PIL import Image
img_pil = Image.open(img_path)
img_pil
input_img = test_transform(img_pil)
input_tensor = input_img.unsqueeze(0).numpy()
input_tensor.shape
(1, 3, 256, 256)
ONNX Runtime预测
注意，输入输出张量的名称需要和 torch.onnx.export 中设置的输入输出名对应
# ONNX Runtime 输入
ort_inputs = {'input': input_tensor}
# ONNX Runtime 输出
pred_logits = ort_session.run(['output'], ort_inputs)[0]
pred_logits = torch.tensor(pred_logits)
import torch.nn.functional as F
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算
pred_softmax.shape
torch.Size([1, 1000])
对预测结果进行柱状图可视化
import matplotlib.pyplot as plt
%matplotlib inline
plt.figure(figsize=(8,4))
x = range(1000)
y = pred_softmax.cpu().detach().numpy()[0]
ax = plt.bar(x, y, alpha=0.5, width=0.3, color='yellow', edgecolor='red', lw=3)
plt.ylim([0, 1.0]) # y轴取值范围
# plt.bar_label(ax, fmt='%.2f', fontsize=15) # 置信度数值
plt.xlabel('Class', fontsize=20)
plt.ylabel('Confidence', fontsize=20)
plt.tick_params(labelsize=16) # 坐标文字大小
plt.title(img_path, fontsize=25)
plt.show()
推理引擎ONNX Runtime部署-ImageNet预训练图像分类模型预测摄像头实时画面
前期准备
#导入包
import onnxruntime
import torch
import pandas as pd
import numpy as np
from PIL import Image, ImageFont, ImageDraw
import matplotlib.pyplot as plt
%matplotlib inline
import os
import numpy as np
import pandas as pd
import cv2 # opencv-python
from tqdm import tqdm # 进度条
import matplotlib.pyplot as plt
%matplotlib inline
import torch
import torch.nn.functional as F
from torchvision import models
# 导入中文字体，指定字号
font = ImageFont.truetype('SimHei.ttf', 32)
# 载入 onnx 模型，获取 ONNX Runtime 推理器
ort_session = onnxruntime.InferenceSession('resnet18.onnx')
# 载入ImageNet 1000图像分类标签
df = pd.read_csv('imagenet_class_index.csv')
idx_to_labels = {}
for idx, row in df.iterrows():
idx_to_labels[row['ID']] = row['Chinese']
# 测试集图像预处理-RCTN：缩放裁剪、转 Tensor、归一化
from torchvision import transforms
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])
])
预测摄像头的一帧画面
获取摄像头的一帧画面
# 导入opencv-python
import cv2
import time
# 获取摄像头，传入0表示获取系统默认摄像头
cap = cv2.VideoCapture(1)
# 打开cap
cap.open(0)
time.sleep(1)
success, img_bgr = cap.read()
# 关闭摄像头
cap.release()
# 关闭图像窗口
cv2.destroyAllWindows()
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) # BGR转RGB
img_pil = Image.fromarray(img_rgb)
img_pil
使用onnx runtime进行预测
# 预处理
input_img = test_transform(img_pil)
input_tensor = input_img.unsqueeze(0).numpy()
# onnx runtime 预测
# onnx runtime 输入
ort_inputs = {'input': input_tensor}
# onnx runtime 输出
pred_logits = ort_session.run(['output'], ort_inputs)[0]
pred_logits = torch.tensor(pred_logits)
import torch.nn.functional as F
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算
## 解析图像分类预测结果
n = 5
top_n = torch.topk(pred_softmax, n) # 取置信度最大的 n 个结果
pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析出类别
confs = top_n[0].cpu().detach().numpy().squeeze() # 解析出置信度
draw = ImageDraw.Draw(img_pil)
# 在图像上写字
for i in range(len(confs)):
pred_class = idx_to_labels[pred_ids[i]]
text = '{:<15} {:>.3f}'.format(pred_class, confs[i])
# 文字坐标，中文字符串，字体，rgba颜色
draw.text((50, 100 + 50 * i), text, font=font, fill=(255, 0, 0, 1))
img = np.array(img_pil) # PIL 转 array
plt.imshow(img)
plt.show()
预测摄像头实时画面
处理单帧画面的函数（中文）
# 处理帧函数
def process_frame(img):
# 记录该帧开始处理的时间
start_time = time.time()
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # BGR转RGB
img_pil = Image.fromarray(img_rgb) # array 转 PIL
## 预处理
input_img = test_transform(img_pil) # 预处理
input_tensor = input_img.unsqueeze(0).numpy()
## onnx runtime 预测
ort_inputs = {'input': input_tensor} # onnx runtime 输入
pred_logits = ort_session.run(['output'], ort_inputs)[0] # onnx runtime 输出
pred_logits = torch.tensor(pred_logits)
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算
## 解析图像分类预测结果
n = 5
top_n = torch.topk(pred_softmax, n) # 取置信度最大的 n 个结果
pred_ids = top_n[1].cpu().detach().numpy().squeeze() # 解析出类别
confs = top_n[0].cpu().detach().numpy().squeeze() # 解析出置信度
draw = ImageDraw.Draw(img_pil)
# 在图像上写字
for i in range(len(confs)):
pred_class = idx_to_labels[pred_ids[i]]
text = '{:<15} {:>.3f}'.format(pred_class, confs[i])
# 文字坐标，中文字符串，字体，rgba颜色
draw.text((50, 100 + 50 * i), text, font=font, fill=(255, 0, 0, 1))
img = np.array(img_pil) # PIL 转 array
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) # RGB转BGR
# 记录该帧处理完毕的时间
end_time = time.time()
# 计算每秒处理图像帧数FPS
FPS = 1/(end_time - start_time)
# 图片，添加的文字，左上角坐标，字体，字体大小，颜色，线宽，线型
img = cv2.putText(img, 'FPS '+str(int(FPS)), (50, 80), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 255), 4, cv2.LINE_AA)
return img
调用摄像头获取每帧
# 调用摄像头逐帧实时处理模板
# 不需修改任何代码，只需修改process_frame函数即可
# 导入opencv-python
import cv2
import time
# 获取摄像头，传入0表示获取系统默认摄像头
cap = cv2.VideoCapture(1)
# 打开cap
cap.open(0)
# 无限循环，直到break被触发
while cap.isOpened():
# 获取画面
success, frame = cap.read()
if not success:
print('Error')
break
## !!!处理帧函数
frame = process_frame(frame)
# 展示处理后的三通道图像
cv2.imshow('my_window',frame)
if cv2.waitKey(1) in [ord('q'),27]: # 按键盘上的q或esc退出（在英文输入法下）
break
# 关闭摄像头
cap.release()
# 关闭图像窗口
cv2.destroyAllWindows()
部署自己训练的图像分类模型
基本流程和上一小节是一样的，我们这里仅展示运用推理引擎ONNX Runtime部署自己训练的水果图像分类模型并预测单张图像。
导出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('checkpoints/fruit30_pytorch_20230123.pth')
model = model.eval().to(device)
Pytorch模型转ONNX模型
x = torch.randn(1, 3, 256, 256).to(device)
with torch.no_grad():
torch.onnx.export(
model, # 要转换的模型
x, # 模型的任意一组输入
'fruit30_resnet18.onnx', # 导出的 ONNX 文件名
opset_version=11, # ONNX 算子集版本
input_names=['input'], # 输入 Tensor 的名称（自己起名字）
output_names=['output'] # 输出 Tensor 的名称（自己起名字）
)
验证onnx模型导出成功
import onnx
# 读取 ONNX 模型
onnx_model = onnx.load('fruit30_resnet18.onnx')
# 检查模型格式是否正确
onnx.checker.check_model(onnx_model)
print('无报错，onnx模型载入成功')
#以可读的形式打印计算图
#print(onnx.helper.printable_graph(onnx_model.graph))
无报错，onnx模型载入成功
和上一小节一样也可以使用Netron对onnx模型可视化
推理引擎ONNX Runtime部署-预测单张图像
import onnxruntime
import numpy as np
import torch
载入 onnx 模型，获取 ONNX Runtime 推理器
ort_session = onnxruntime.InferenceSession('fruit30_resnet18.onnx')
预处理
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])
])
img_path = 'test_img/watermelon1.jpg'
# 用 pillow 载入
from PIL import Image
img_pil = Image.open(img_path)
img_pil
input_img = test_transform(img_pil)
input_tensor = input_img.unsqueeze(0).numpy()
input_tensor.shape
(1, 3, 256, 256)
ONNX Runtime预测
注意，输入输出张量的名称需要和 torch.onnx.export 中设置的输入输出名对应
# ONNX Runtime 输入
ort_inputs = {'input': input_tensor}
# ONNX Runtime 输出
pred_logits = ort_session.run(['output'], ort_inputs)[0]
pred_logits = torch.tensor(pred_logits)
import torch.nn.functional as F
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算
pred_softmax.shape
torch.Size([1, 30])
解析预测结果
#设置matplotlib中文字体
# windows操作系统
plt.rcParams['font.sans-serif']=['SimHei'] # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus']=False # 用来正常显示负号
#载入类别和对应 ID
idx_to_labels = np.load('idx_to_labels.npy', allow_pickle=True).item()
绘制预测结果柱状图
import matplotlib.pyplot as plt
%matplotlib inline
plt.figure(figsize=(22, 10))
x = idx_to_labels.values()
y = pred_softmax.cpu().detach().numpy()[0] * 100
width = 0.45 # 柱状图宽度
ax = plt.bar(x, y, width)
plt.bar_label(ax, fmt='%.2f', fontsize=15) # 置信度数值
plt.tick_params(labelsize=20) # 设置坐标文字大小
plt.title(img_path, fontsize=30)
plt.xticks(rotation=45) # 横轴文字旋转
plt.xlabel('类别', fontsize=20)
plt.ylabel('置信度', fontsize=20)
plt.show()
总结
本文主要讲述了ONNX-ONNX Runtime部署流程，首先将训练好的Pytorch模型转ONNX模型，这样我们就可以将ONNX模型在任何安装了ONNX Runtime环境的机器上进行运行，进行单张图片的预测、调用摄像头进行实时画面的预测等。使用ONNX我们可以让模型在不同框架之间进行迁移，方便我们低成本的将模型部署到移动设备中去。