将模型文件转换为 ONNX 格式并使用 ONNX Runtime 进行推理

news2024/11/26 16:46:17

如何将模型文件转换为 ONNX 格式并使用 ONNX Runtime 进行推理

ONNX（Open Neural Network Exchange）是一种开放格式，用于表示机器学习模型。它允许不同框架之间的模型互操作性，使得模型可以在不同的推理引擎中运行。本文将详细介绍如何将模型文件转换为 ONNX 格式，并使用 ONNX Runtime 进行推理。

1. 环境搭建

1.1 安装必要的库

首先，确保您的系统已经安装了 Python 和相关的库。这里我们将使用 PyTorch 作为示例，但 ONNX 支持多种框架，如 TensorFlow、Keras 等。

pip install torch onnx onnxruntime

1.2 检查 CUDA 支持

如果您的系统有 GPU，确保 CUDA 和 cuDNN 已正确安装，并且 PyTorch 能够检测到 GPU。

import torch

print(torch.cuda.is_available())  # 应输出 True
print(torch.cuda.device_count())  # 应输出 GPU 的数量

2. 将模型转换为 ONNX 格式

2.1 加载预训练模型

假设我们有一个使用 PyTorch 训练的模型，首先加载该模型。

import torch
from torchvision.models import resnet18

# 加载预训练模型
model = resnet18(pretrained=True)
model.eval()  # 设置模型为评估模式

2.2 定义输入张量

定义一个输入张量，用于模型的前向传播。输入张量的形状应与模型的输入要求一致。

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

2.3 导出模型为 ONNX 格式

使用 torch.onnx.export 函数将模型导出为 ONNX 格式。

import torch.onnx

# 导出模型为 ONNX 格式
torch.onnx.export(
    model,                  # 模型对象
    dummy_input,            # 输入张量
    "resnet18.onnx",        # 输出文件名
    export_params=True,     # 是否导出模型参数
    opset_version=11,       # ONNX 操作集版本
    do_constant_folding=True,  # 是否执行常量折叠优化
    input_names=['input'],  # 输入节点名称
    output_names=['output'],  # 输出节点名称
    dynamic_axes={'input': {0: 'batch_size'},    # 动态轴
                  'output': {0: 'batch_size'}}
)

3. 使用 ONNX Runtime 进行推理

3.1 安装 ONNX Runtime

确保已经安装了 ONNX Runtime。

pip install onnxruntime

3.2 加载 ONNX 模型

使用 ONNX Runtime 加载导出的 ONNX 模型。

import onnxruntime as ort
import numpy as np

# 加载 ONNX 模型
ort_session = ort.InferenceSession("resnet18.onnx")

3.3 准备输入数据

准备输入数据，将其转换为 NumPy 数组。

# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

3.4 进行推理

使用 ONNX Runtime 进行推理，并获取输出结果。

# 进行推理
outputs = ort_session.run(None, {'input': input_data})

# 获取输出结果
output = outputs[0]
print(output.shape)  # 输出形状应为 (1, 1000)，对应 1000 个类别的概率分布

4. 性能优化

4.1 使用 GPU 进行推理

如果您的系统有 GPU，可以配置 ONNX Runtime 使用 GPU 进行推理。

# 配置使用 GPU
ort_session = ort.InferenceSession("resnet18.onnx", providers=['CUDAExecutionProvider'])

# 进行推理
outputs = ort_session.run(None, {'input': input_data})

# 获取输出结果
output = outputs[0]
print(output.shape)

4.2 动态批处理

ONNX Runtime 支持动态批处理，可以根据实际需求调整批处理大小。

# 动态批处理
batch_size = 4
input_data = np.random.randn(batch_size, 3, 224, 224).astype(np.float32)

# 进行推理
outputs = ort_session.run(None, {'input': input_data})

# 获取输出结果
output = outputs[0]
print(output.shape)  # 输出形状应为 (4, 1000)