MMDeploy项目中RTMO模型在TensorRT部署时的矩形输入问题解析

问题背景

在MMDeploy项目中将RTMO姿态估计模型部署到TensorRT时，开发者发现了一个有趣的现象：当使用正方形输入尺寸(1×3×640×640)时，TensorRT模型能够正常工作，输出结果与PyTorch模型基本一致；但当改为矩形输入尺寸(1×3×1280×640)时，虽然PyTorch模型仍能正常工作，TensorRT模型输出的关键点位置却出现了明显偏差。

现象表现

从实际运行结果对比图中可以观察到：

• PyTorch模型在矩形输入下输出正常，关键点位置准确
• TensorRT模型在相同输入下，虽然整体姿态相似，但关键点位置存在明显偏移

问题根源分析

经过深入排查，发现问题出在RTMO头部(rtmo_head)的重写器(rewriter)实现上。具体来说：

1. RTMO头部在部署模式下需要两个关键变量：self.flatten_priors和self.flatten_stride
2. 这两个变量是在switch_to_deploy函数中预定义的（位于mmpose/models/heads/hybrid_heads/rtmo_head.py）
3. 原实现中特征图生成时使用了固定的高宽比例，没有正确处理矩形输入的情况

解决方案

问题的核心在于特征图生成时的尺寸处理不当。在PyTorch张量中，尺寸顺序是(1,1,height,width)，而输入尺寸参数input_size的顺序是(width,height)。正确的修改方式应该是：

featuremaps.append(torch.rand(1, 1, input_size[1] // s, input_size[0] // s))

这一修改确保了：

1. 正确解析输入尺寸的宽高顺序
2. 特征图生成时保持正确的宽高比例
3. 适应各种矩形输入尺寸的情况

技术启示

这个案例给我们几点重要启示：

1. 尺寸顺序一致性：在计算机视觉和深度学习领域，不同框架和接口对尺寸顺序的定义可能不同，开发时需要特别注意
2. 部署验证：模型在训练框架和部署框架中的行为可能存在差异，需要进行全面的验证测试
3. 矩形输入支持：许多视觉模型最初设计时考虑正方形输入，但在实际应用中矩形输入更为常见，部署时需要确保支持

总结

通过分析RTMO模型在TensorRT部署时出现的矩形输入问题，我们不仅找到了具体解决方案，也加深了对模型部署过程中尺寸处理重要性的理解。这类问题的解决往往需要对训练框架和部署框架都有深入理解，才能准确定位问题根源。

对于开发者而言，在模型部署过程中，应当特别注意输入输出尺寸的匹配问题，特别是在处理非正方形输入时，需要全面验证模型在各个组件中的尺寸处理逻辑。