mmdetection3d模型导出指南：ONNX格式转换与验证

2026-02-05 05:48:49作者：曹令琨Iris

1. 引言

在计算机视觉领域，3D目标检测（3D Object Detection）技术在自动驾驶、机器人导航等领域发挥着关键作用。mmdetection3d作为OpenMMLab推出的下一代通用3D目标检测平台，提供了丰富的模型和算法。将训练好的模型导出为ONNX（Open Neural Network Exchange）格式，可以实现跨框架部署、模型优化和性能评估等功能。本文将详细介绍如何使用mmdetection3d将模型导出为ONNX格式，并进行验证。

2. 环境准备

2.1 安装依赖

在开始模型导出之前，需要确保已安装以下依赖库：

# 克隆仓库
git clone https://gitcode.com/gh_mirrors/mm/mmdetection3d.git
cd mmdetection3d

# 安装依赖
pip install -r requirements.txt
pip install onnx onnxruntime

2.2 准备模型和配置文件

mmdetection3d提供了多种预训练模型，可通过官方模型库下载。例如，我们以PointPillars模型为例进行导出，其配置文件位于configs/pointpillars/pointpillars_hv_secfpn_8xb6-160e_kitti-3d-car.py。

3. 模型导出流程

3.1 导出工具介绍

mmdetection3d虽然没有专门的ONNX导出脚本，但可以通过修改现有代码或编写自定义脚本来实现。主要思路是使用PyTorch的torch.onnx.export函数将模型导出为ONNX格式。

3.2 编写导出脚本

创建export_onnx.py脚本，内容如下：

import torch
from mmdet3d.apis import init_model
import onnx
import onnxruntime as ort
import numpy as np

def export_model_to_onnx(config_path, checkpoint_path, output_path, input_shape=(1, 4, 16384)):
    # 初始化模型
    model = init_model(config_path, checkpoint_path, device='cpu')
    model.eval()

    # 创建输入张量
    input_tensor = torch.randn(input_shape)

    # 导出ONNX模型
    torch.onnx.export(
        model,
        input_tensor,
        output_path,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}},
        opset_version=11
    )

    print(f"ONNX模型已导出至 {output_path}")

def verify_onnx_model(onnx_path, input_shape=(1, 4, 16384)):
    # 加载ONNX模型
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)

    # 创建ONNX Runtime会话
    ort_session = ort.InferenceSession(onnx_path)

    # 生成随机输入
    input_data = np.random.randn(*input_shape).astype(np.float32)

    # 推理
    inputs = {ort_session.get_inputs()[0].name: input_data}
    outputs = ort_session.run(None, inputs)

    print(f"ONNX模型验证成功，输出形状: {outputs[0].shape}")

if __name__ == "__main__":
    config_path = "configs/pointpillars/pointpillars_hv_secfpn_8xb6-160e_kitti-3d-car.py"
    checkpoint_path = "checkpoints/pointpillars_hv_secfpn_8xb6-160e_kitti-3d-car_20220331_134606-d42d15ed.pth"
    output_path = "pointpillars.onnx"

    export_model_to_onnx(config_path, checkpoint_path, output_path)
    verify_onnx_model(output_path)

3.3 执行导出脚本

python export_onnx.py

4. 模型验证

4.1 检查ONNX模型结构

使用Netron工具可视化ONNX模型结构，确保网络层和输入输出节点正确。

pip install netron
netron pointpillars.onnx

4.2 精度验证

比较PyTorch模型和ONNX模型在相同输入下的输出结果，确保精度损失在可接受范围内。

# 在export_onnx.py中添加精度验证代码
def compare_outputs(pytorch_model, onnx_session, input_data):
    # PyTorch推理
    with torch.no_grad():
        pytorch_output = pytorch_model(torch.tensor(input_data)).numpy()

    # ONNX推理
    onnx_output = onnx_session.run(None, {onnx_session.get_inputs()[0].name: input_data})[0]

    # 计算误差
    mse = np.mean((pytorch_output - onnx_output) ** 2)
    print(f"PyTorch与ONNX输出的MSE: {mse}")

    return mse < 1e-5  # 设定误差阈值

5. 常见问题与解决方案

5.1 动态输入形状问题

mmdetection3d中的部分模型（如基于点云的模型）输入点云数量不固定，导出ONNX时需指定动态轴：

dynamic_axes={'input': {0: 'batch_size', 2: 'num_points'}, 'output': {0: 'batch_size'}}

5.2 不支持的算子

若导出过程中出现不支持的算子错误，可尝试以下解决方案：

更新PyTorch和ONNX版本
使用torch.onnx.export的operator_export_type参数指定算子导出类型
自定义算子或修改模型结构

5.3 模型优化

导出ONNX模型后，可使用ONNX Runtime进行优化：

from onnxruntime.quantization import quantize_dynamic, QuantType

quantize_dynamic(
    'pointpillars.onnx',
    'pointpillars_quantized.onnx',
    weight_type=QuantType.QUInt8
)

6. 总结与展望

本文详细介绍了mmdetection3d模型导出为ONNX格式的流程，包括环境准备、导出脚本编写、模型验证等步骤。通过将模型导出为ONNX格式，可以实现跨框架部署和模型优化，为3D目标检测模型的实际应用提供了便利。未来，随着mmdetection3d的不断更新，官方可能会提供更完善的ONNX导出工具，进一步简化导出流程。