ONNX-ONNX Runtime部署
1. 部署ImageNet預(yù)訓(xùn)練圖像分類模型
#安裝配置環(huán)境
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
導(dǎo)出ONNX模型(把原生pytorch訓(xùn)練得到的圖像分類模型铛绰,導(dǎo)出為ONNX格式染坯,用于后續(xù)在ONNX Runtime推理引擎上部署)
#導(dǎo)入工具包
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預(yù)訓(xùn)練圖像分類模型
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模型轉(zhuǎn)為ONNX模型
x = torch.randn(1, 3, 256, 256).to(device)
with torch.no_grad():
torch.onnx.export(
model, # 要轉(zhuǎn)換的模型
x, # 模型的任意一組輸入
'resnet18.onnx', # 導(dǎo)出的 ONNX 文件名
opset_version=11, # ONNX 算子集版本
input_names=['input'], # 輸入 Tensor 的名稱(自己起名字)
output_names=['output'] # 輸出 Tensor 的名稱(自己起名字)
)
#驗(yàn)證ONNX模型導(dǎo)出成功
import onnx
# 讀取 ONNX 模型
onnx_model = onnx.load('resnet18.onnx')
# 檢查模型格式是否正確
onnx.checker.check_model(onnx_model)
print('無報(bào)錯(cuò)橙数,onnx模型載入成功')
#以可讀的形式打印計(jì)算圖
print(onnx.helper.printable_graph(onnx_model.graph))
## 使用Netron對onnx模型可視化
https://netron.app
resnet18.onnx
使用ONNX尽楔,讀取onnx格式的模型文件,對單張圖像文件進(jìn)行預(yù)測
#導(dǎo)入工具包
import onnxruntime
import numpy as np
import torch
#載入onnx模型
ort_session = onnxruntime.InferenceSession('resnet18.onnx')
#構(gòu)造輸入拱她,獲取輸出結(jié)果
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
#預(yù)處理
from torchvision import transforms
# 測試集圖像預(yù)處理-RCTN:縮放裁剪、轉(zhuǎn) 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
#運(yùn)行預(yù)處理
input_img = test_transform(img_pil)
input_img.shape
input_tensor = input_img.unsqueeze(0).numpy()
input_tensor.shape
#使用ONNX預(yù)測
# 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 分?jǐn)?shù)做 softmax 運(yùn)算
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) # 置信度數(shù)值
plt.xlabel('Class', fontsize=20)
plt.ylabel('Confidence', fontsize=20)
plt.tick_params(labelsize=16) # 坐標(biāo)文字大小
plt.title(img_path, fontsize=25)
plt.show()
2. 部署自己訓(xùn)練的水果圖像分類模型
導(dǎo)出ONNX模型
#導(dǎo)入工具包
import torch
from torchvision import models
# 有 GPU 就用 GPU秉沼,沒有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device', device)
#導(dǎo)入訓(xùn)練好的模型
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模型轉(zhuǎn)為ONNX模型
x = torch.randn(1, 3, 256, 256).to(device)
with torch.no_grad():
torch.onnx.export(
model, # 要轉(zhuǎn)換的模型
x, # 模型的任意一組輸入
'fruit30_resnet18.onnx', # 導(dǎo)出的 ONNX 文件名
opset_version=11, # ONNX 算子集版本
input_names=['input'], # 輸入 Tensor 的名稱(自己起名字)
output_names=['output'] # 輸出 Tensor 的名稱(自己起名字)
)
#驗(yàn)證ONNX模型導(dǎo)出成功import onnx
# 讀取 ONNX 模型
onnx_model = onnx.load('fruit30_resnet18.onnx')
# 檢查模型格式是否正確
onnx.checker.check_model(onnx_model)
print('無報(bào)錯(cuò),onnx模型載入成功')
#以可讀的形式打印計(jì)算圖
print(onnx.helper.printable_graph(onnx_model.graph))
### 使用Netron對onnx模型可視化 https://netron.app
fruit30_resnet18.onnx
使用ONNX矿酵,讀取onnx格式的模型文件唬复,預(yù)測單張圖像
#導(dǎo)入工具包
import onnxruntime
import numpy as np
import torch
#載入onnx模型
ort_session = onnxruntime.InferenceSession('fruit30_resnet18.onnx')
#構(gòu)造輸入,獲取輸出結(jié)果
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
#預(yù)處理
from torchvision import transforms
# 測試集圖像預(yù)處理-RCTN:縮放裁剪全肮、轉(zhuǎn) 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預(yù)測
# 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 分?jǐn)?shù)做 softmax 運(yùn)算
pred_softmax.shape
#解析預(yù)測結(jié)果
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) # 置信度數(shù)值
plt.tick_params(labelsize=20) # 設(shè)置坐標(biāo)文字大小
plt.title(img_path, fontsize=30)
plt.xticks(rotation=45) # 橫軸文字旋轉(zhuǎn)
plt.xlabel('類別', fontsize=20)
plt.ylabel('置信度', fontsize=20)
plt.show()