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}")
