onnx 图像分类

参考文章：

【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.shape

?Pytorch模型转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()