ONNX-ONNX Runtime部署
1. 部署ImageNet预训练图像分类模型
#安装配置环境
pip3 install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113
pip install numpy pandas matplotlib tqdm opencv-python pillow onnx onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple
导出ONNX模型（把原生pytorch训练得到的图像分类模型，导出为ONNX格式，用于后续在ONNX Runtime推理引擎上部署）
#导入工具包
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)
x = torch.randn(1, 3, 256, 256).to(device)
output=model(x)
output.shape
#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))
## 使用Netron对onnx模型可视化
https://netron.app
resnet18.onnx
使用ONNX，读取onnx格式的模型文件，对单张图像文件进行预测
#导入工具包
import onnxruntime
import numpy as np
import torch
#载入onnx模型
ort_session = onnxruntime.InferenceSession('resnet18.onnx')
#构造输入，获取输出结果
x = torch.randn(1, 3, 256, 256).numpy()
x.shape
# onnx runtime 输入
ort_inputs = {'input': x}
# onnx runtime 输出
ort_output = ort_session.run(['output'], ort_inputs)[0]
ort_output.shape
#预处理
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_img.shape
input_tensor = input_img.unsqueeze(0).numpy()
input_tensor.shape
#使用ONNX预测
# 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
import torch.nn.functional as F
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算
pred_softmax.shape
#柱状图可视化
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()
2. 部署自己训练的水果图像分类模型
导出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_20220814.pth')
model = model.eval().to(device)
x = torch.randn(1, 3, 256, 256).to(device)
output = model(x)
output.shape
#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))
### 使用Netron对onnx模型可视化 https://netron.app
fruit30_resnet18.onnx
使用ONNX，读取onnx格式的模型文件，预测单张图像
#导入工具包
import onnxruntime
import numpy as np
import torch
#载入onnx模型
ort_session = onnxruntime.InferenceSession('fruit30_resnet18.onnx')
#构造输入，获取输出结果
x = torch.randn(1, 3, 256, 256).numpy()
x.shape
# onnx runtime 输入
ort_inputs = {'input': x}
# onnx runtime 输出
ort_output = ort_session.run(['output'], ort_inputs)[0]
ort_output.shape
#预处理
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_img.shape
input_tensor = input_img.unsqueeze(0).numpy()
input_tensor.shape
#ONNX预测
# 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
import torch.nn.functional as F
pred_softmax = F.softmax(pred_logits, dim=1) # 对 logit 分数做 softmax 运算
pred_softmax.shape
#解析预测结果
idx_to_labels = np.load('idx_to_labels.npy', allow_pickle=True).item()
idx_to_labels
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()
