pytorch模型转ONNX、并进行比较推理

pytorch模型转ONNX

概述

ONNX（Open Neural Network Exchange）是一种开放式的深度学习模型交换格式，旨在促进不同深度学习框架之间的互操作性。通过将深度学习模型转换为ONNX格式，可以将其从一个深度学习框架移植到另一个框架中，而无需重新训练模型或手动重新实现模型结构。

ONNX是由微软和Facebook于2017年联合发布的，目前得到了众多公司和社区的支持和贡献。ONNX格式支持多种深度学习模型类型，包括卷积神经网络（CNN）、循环神经网络（RNN）和变换器（Transformer）等。

在使用ONNX时，通常需要两个步骤：模型的导出和模型的导入。首先，将模型从一个深度学习框架中导出到ONNX格式。导出后，可以使用ONNX运行时（ONNX Runtime）或其他支持ONNX格式的深度学习框架将模型导入到所需的平台中进行推理。

在导出模型时，需要注意确保模型的兼容性。由于每个框架实现深度学习算法的方式不同，因此在转换模型时可能会出现一些限制或不兼容性。此外，由于ONNX不支持某些框架特有的操作或层，因此在导出模型时可能需要进行一些自定义操作或调整模型结构，以便使模型兼容ONNX格式。

总的来说，ONNX是一个非常有用的工具，可以简化深度学习模型的转换和部署过程，促进不同框架之间的互操作性和合作。

PT转ONNX

import torch
import onnx
import onnxruntime
from efficient_net import build_model


# Step 1: 加载 PyTorch 模型
# model = EfficientNet.from_pretrained('efficientnet-b0')
model = build_model(model_name="efficientnet-b2",weights_path="ef_classify/models/best.pt", num_classes=2)
model.eval()

# Step 2: 将模型转换为 ONNX 格式（包含动态批次）
output_path = "efficientnet-b2.onnx"  # ONNX 模型保存路径

# 导出 ONNX 模型
dummy_input = torch.randn(1, 3, 260, 260)
dynamic_axes = {'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
torch.onnx.export(model, dummy_input, output_path, verbose=True, input_names=['input'], output_names=['output'],
                  dynamic_axes=dynamic_axes)

比较PT和ONNX

要比较转换后的ONNX模型和原始的PyTorch模型之间的差异，可以使用ONNX Runtime工具包。

1. 加载模型：在使用ONNX Runtime进行模型比较之前，您需要加载PyTorch模型和转换后的ONNX模型。您可以使用PyTorch或ONNX Runtime加载模型。

# 加载 PyTorch 模型
import torch
from efficient_net import build_model
# model_pt = torch.load('model.pt')
model_pt = build_model(model_name="efficientnet-b2",weights_path="ef_classify/models/best.pt", num_classes=2)

# 加载转换后的 ONNX 模型
import onnxruntime as ort
model_onnx = ort.InferenceSession('efficientnet-b2.onnx',providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

2. 输入数据：为了比较模型，需要提供相同的输入数据。可以生成一些随机数据并将其提供给模型进行比较。

import numpy as np
input_shape = (1, 3, 260, 260)
input_data = np.random.random(input_shape).astype(np.float32)

3. 推理和比较：使用PyTorch和ONNX Runtime对相同的输入数据进行推理并比较输出结果。如果输出相同，则说明两个模型在这些输入数据上是等价的。

# 使用 PyTorch 进行推理
model_pt.eval()
with torch.no_grad():
    output_pt = model_pt(torch.tensor(input_data))

# 使用 ONNX Runtime 进行推理
output_onnx = model_onnx.run(None, {'input': input_data})[0]

# 比较输出结果
if np.allclose(output_pt.numpy(), output_onnx, rtol=1e-3, atol=1e-5):
    print("模型输出相同！")
else:
    print("模型输出不同！")

使用这种方法，可以比较转换后的ONNX模型和原始的PyTorch模型在相同输入数据上的输出结果是否相同，从而评估模型的差距。

ONNX批量推理图片

在下面的代码中，我们使用了一个批处理的方式来进行推理。我们首先获取了所有图像文件的路径，然后将其分成了多个大小为batch_size的批次进行推理。在每个批次中，我们将多张图像一次性输入模型进行推理，并将输出结果保存在outputs列表中。最后，我们将所有输出结果合并并输出每个图像的预测结果。

import os
import cv2
import numpy as np
import onnxruntime as ort

# 加载ONNX模型
model_path = "efficientnet-b2.onnx"
session =  ort.InferenceSession(model_path,providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

input_name = session.get_inputs()[0].name

# 获取所有图像文件的路径
images_path = "ef_classify/dataset/dog"
image_paths = [os.path.join(images_path, f) for f in os.listdir(images_path) if os.path.isfile(os.path.join(images_path, f))]


# 批量读取、预处理和推理图像
batch_size = 8
outputs = []
for i in range(0, len(image_paths), batch_size):
    batch_images = []
    for image_path in image_paths[i:i+batch_size]:
        # 读取图像并进行预处理
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = cv2.resize(image, (260, 260))
        img = image.astype(np.float32)
        
        img /= 255.0
        img_mean = (0.485,0.456,0.406)
        img_std = (0.229,0.224,0.225)
        img = img - img_mean
        img = (img/img_std)

        image = img.transpose((2, 0, 1))
        image = np.expand_dims(image, axis=0).astype(np.float32)
        batch_images.append(image)

    # 进行推理并保存输出结果
    batch_images = np.concatenate(batch_images, axis=0)
    batch_outputs = session.run(None, {input_name: batch_images})
    outputs.append(batch_outputs[0])

# 将所有输出结果合并并输出每个图像的预测结果
outputs = np.concatenate(outputs, axis=0)
for i, image_path in enumerate(image_paths):
    predicted_class = np.argmax(outputs[i])
    print(f"Image: {image_path} - Predicted Class: {predicted_class}")