onnxruntime 推理分类模型 以resnet为例

onnxruntime 推理分类模型 以resnet为例

分类模型的预处理（1） Resize 和 CenterCrop
分类模型的预处理（2）ToTensor 和 Normalize
分类模型 Top 1和Top 5 指标说明
分类模型 Top 1和Top 5 指标实现

softmax 函数的多种实现方式 包括纯C语言、C++版本、Eigen版本等
Softmax用于二分类时会退化为Sigmoid

import torch
import torchvision.models as models
import torch.onnx
import onnx
import onnxruntime as ort
import numpy as np
from PIL import Image
from torchvision import transforms

# 1. 定义并加载预训练的ResNet模型
model = models.resnet18(pretrained=True)
model.eval()

# 2. 定义输入
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出为ONNX格式
onnx_model_path = "resnet18.onnx"
torch.onnx.export(model, dummy_input, onnx_model_path, verbose=True)

# 4. 加载ONNX模型
onnx_model = onnx.load(onnx_model_path)


# 5. 使用ONNX Runtime进行推断
# 加载ONNX模型
ort_session = ort.InferenceSession(onnx_model_path)

# 图像预处理
def preprocess_image(image_path):
    input_image = Image.open(image_path)
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    input_tensor = preprocess(input_image)
    input_batch = input_tensor.unsqueeze(0)  # 创建batch维度
    return input_batch.numpy()

# 推断函数
def infer(image_path):
    input_batch = preprocess_image(image_path)
    ort_inputs = {ort_session.get_inputs()[0].name: input_batch}
    ort_outs = ort_session.run(None, ort_inputs)
    return ort_outs

# Softmax函数
def softmax(x):
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum(axis=1, keepdims=True)

# 示例图像路径
image_path = "1.jpg"

# 进行推断
outputs = infer(image_path)

# 输出的结构是一个列表，其中第一个元素是模型的输出，形状为 (1, num_classes)
print(f"outputs 的结构: {type(outputs)}, 长度: {len(outputs)}")
print(f"outputs[0] 的结构: {type(outputs[0])}, 形状: {outputs[0].shape}")

# 应用softmax
probs = softmax(outputs[0])

# 获取预测的类别索引和置信度
predicted_class = np.argmax(probs)
confidence = probs[0][predicted_class]

print(f"预测的类别索引：{predicted_class}")
print(f"预测的置信度：{confidence}")

# 可以加载类别名称文件进行实际的类别名称映射
# 例如：imagenet_labels = ["class0", "class1", ..., "class999"]
# print(f"预测的类别名称：{imagenet_labels[predicted_class]}")

在 transforms.Compose 中，变换的执行顺序是按列出的顺序进行的。
transforms.Resize(256)：首先将图像的短边调整到 256 像素，长边按比例缩放。
transforms.CenterCrop(224)：然后从图像中心裁剪出 224x224 的区域。
transforms.ToTensor()：接着将图像转换为张量，并将像素值归一化到 [0, 1] 之间。
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])：最后使用给定的均值和标准差对图像进行归一化。
所以执行顺序是先 Resize，再 CenterCrop，接着 ToTensor，最后 Normalize。