PyTorch转ONNX时动态输入尺寸的实战指南：以RetinaFace多输出模型为例

当我们将PyTorch模型转换为ONNX格式时，动态输入尺寸的设置往往成为工程师们头疼的问题。特别是像RetinaFace这样的人脸检测模型，不仅输入尺寸需要动态调整，多个输出（如边界框、关键点和置信度）的尺寸也需要灵活处理。本文将深入探讨如何正确配置 dynamic_axes 参数，解决实际部署中遇到的尺寸不匹配问题。

1. 理解动态输入尺寸的核心概念

动态输入尺寸允许模型在推理时接受不同尺寸的输入，这在计算机视觉任务中尤为重要。想象一下，在实际应用中，我们可能需要处理各种分辨率的图像，从低分辨率的监控摄像头到高分辨率的专业摄影设备。

PyTorch的 torch.onnx.export 函数提供了 dynamic_axes 参数来实现这一功能。这个参数本质上是一个字典，用于指定哪些维度应该是动态的。例如：

dynamic_axes = {
    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output': {0: 'batch_size'}
}

这里有几个关键点需要注意：

• 维度索引 ：PyTorch使用从0开始的索引来表示张量的各个维度
• 命名约定 ：可以为动态维度指定有意义的名称，便于后续调试
• 输入输出映射 ：需要同时为输入和输出指定动态维度

提示：虽然ONNX支持动态形状，但实际支持程度取决于具体运行时环境。onnxruntime通常对动态形状有较好的支持。

2. RetinaFace模型的多输出挑战

RetinaFace作为先进的人脸检测模型，其输出结构相对复杂。典型的输出包括：

1. 边界框 （Bounding Boxes）：形状为[N, M, 4]
2. 关键点 （Landmarks）：形状为[N, M, 10]（5个关键点，每个点有x,y坐标）
3. 置信度 （Confidence Scores）：形状为[N, M, 2]

当输入图像的尺寸变化时，M（检测到的特征图数量）也会随之变化。这就是为什么我们需要特别关注多输出情况下的动态尺寸配置。

2.1 常见错误模式分析

许多工程师在初次尝试时会遇到类似下面的警告：

[W:onnxruntime:, execution_frame.cc:721 VerifyOutputSizes] 
Expected shape from model of {1,15162,2} does not match 
actual shape of {1,15700,2} for output 837

这种警告表明，虽然模型理论上支持动态尺寸，但实际配置可能存在问题。主要原因包括：

• 没有为所有输出正确指定动态维度
• 动态维度的索引指定错误
• 输出名称与模型定义不匹配

3. 正确配置多输出模型的动态尺寸

针对RetinaFace这样的多输出模型，我们需要更精细地配置 dynamic_axes 参数。以下是一个完整的示例：

dynamic_axes = {
    'input': {
        0: 'batch_size',
        2: 'height', 
        3: 'width'
    },
    'boxes': {
        0: 'batch_size',
        1: 'num_detections'
    },
    'landmarks': {
        0: 'batch_size',
        1: 'num_detections'
    },
    'scores': {
        0: 'batch_size',
        1: 'num_detections'
    }
}

torch.onnx.export(
    model,
    dummy_input,
    'retinaface.onnx',
    export_params=True,
    opset_version=12,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['boxes', 'landmarks', 'scores'],
    dynamic_axes=dynamic_axes
)

3.1 关键配置要点

1. 输入配置 ：

   • 通常需要将batch维度（0）设为动态
   • 对于图像输入，高度和宽度维度（2和3）也应设为动态

2. 输出配置 ：

   • 所有输出的batch维度应与输入一致
   • 检测数量维度（通常是1）必须设为动态
   • 固定维度（如边界框的4个坐标）不应包含在动态配置中

3. 版本兼容性 ：

   • 建议使用opset_version 11或更高版本
   • 较旧版本的ONNX对动态形状支持有限

4. 验证与调试技巧

导出ONNX模型后，我们需要验证动态尺寸是否按预期工作。以下是一些实用的验证方法：

4.1 使用onnxruntime进行验证

import onnxruntime as ort

# 创建推理会话
sess = ort.InferenceSession('retinaface.onnx')

# 准备不同尺寸的输入
input1 = torch.randn(1, 3, 640, 480)  # 480p
input2 = torch.randn(1, 3, 1280, 720) # 720p

# 运行推理
outputs1 = sess.run(None, {'input': input1.numpy()})
outputs2 = sess.run(None, {'input': input2.numpy()})

# 检查输出形状
print(f"480p输入下的输出形状: {[o.shape for o in outputs1]}")
print(f"720p输入下的输出形状: {[o.shape for o in outputs2]}")

4.2 常见问题排查表

问题现象	可能原因	解决方案
运行时形状不匹配警告	动态维度配置不完整	检查是否为所有变化的维度配置了动态轴
推理结果异常	某些操作不支持动态形状	检查模型中的reshape、view等操作
导出失败	使用了不支持的opset版本	尝试更新到更高版本的opset
性能下降	动态形状导致优化受限	考虑使用半动态配置（如仅动态batch）

4.3 高级调试技巧

对于复杂模型，可以使用ONNX的shape inference功能来验证动态形状：

import onnx

model = onnx.load('retinaface.onnx')
onnx.checker.check_model(model)
inferred_model = onnx.shape_inference.infer_shapes(model)
print(inferred_model.graph.value_info)

5. 生产环境最佳实践

在实际部署多输出ONNX模型时，除了正确配置动态尺寸外，还需要考虑以下因素：

1. 内存管理 ：

   • 动态尺寸可能导致内存使用波动
   • 需要设置合理的内存预算和回退机制

2. 性能优化 ：

   • 对于常用尺寸，可以预先编译优化版本
   • 考虑使用ONNX Runtime的IOBinding功能减少数据拷贝

3. 版本控制 ：

   • 记录使用的PyTorch和ONNX版本
   • 为不同硬件平台保存特定优化版本

4. 监控与日志 ：

   • 记录实际运行的输入输出尺寸分布
   • 设置警报应对异常尺寸情况

# 生产环境推荐的导出参数
torch.onnx.export(
    model,
    dummy_input,
    'retinaface_prod.onnx',
    export_params=True,
    opset_version=12,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['boxes', 'landmarks', 'scores'],
    dynamic_axes=dynamic_axes,
    training=torch.onnx.TrainingMode.EVAL,
    operator_export_type=torch.onnx.OperatorExportTypes.ONNX,
    verbose=False
)

在实际项目中，我们发现合理配置动态尺寸可以显著提高模型的部署灵活性，特别是在需要处理多种输入源的应用场景中。对于RetinaFace这样的人脸检测模型，确保所有输出都正确配置了动态维度是避免运行时问题的关键。