前言

最近在 MMPose 上看到了一个效果比较好的模型 RTMO，想通过调试分析 MMPose 代码把 RTMO 模型导出，并在 tensorRT 上推理得到结果，这篇文章主要分析 RTMO 模型的 ONNX 导出以及解决导出过程中遇到的各种问题。若有问题欢迎各位看官批评指正😄

Paper：RTMO: Towards High-Performance One-Stage Real-Time Multi-Person Pose Estimation

Github：https://github.com/open-mmlab/mmpose、https://github.com/open-mmlab/mmdeploy

1. 概述

1.1 MMPopse

以下内容均 Copy 自：https://mmpose.readthedocs.io/en/latest/

MMPose 是一款基于 Pytorch 的姿态点估计开源工具箱，是 OpenMMLab 项目的成员之一，包含了丰富的 2D 多人姿态估计、2D 手部姿态估计、2D 人脸关键点检测、133 个关键点全身人体姿态估计、动物关键点检测、服饰关键点检测等算法以及相关的组件和模块。

MMPose 由以下 8 个部分组成：

• apis 提供用于模型推理的高级 API

• structures 提供 bbox、keypoint 和 PoseDataSample 等数据结构

• datasets 支持用于姿态估计的各种数据集

  • transforms 包含各种数据增强变换

• codecs 提供姿态编解码器：编码器用于将姿态信息（通常为关键点坐标）编码为模型学习目标（如热力图），解码器则用于将模型输出解码为姿态估计结果

• models 以模块化结构提供了姿态估计模型的各类组件

  • pose_estimators 定义了所有姿态估计模型类
  • data_preprocessors 用于预处理模型的输入数据
  • backbones 包含各种骨干网络
  • necks 包含各种模型颈部组件
  • heads 包含各种模型头部
  • losses 包含各种损失函数

• engine 包含与姿态估计任务相关的运行时组件

  • hooks 提供运行时的各种钩子

• evaluation 提供各种评估模型性能的指标

• visualization 用于可视化关键点骨架和热力图等信息

1.2 MMDeploy

以下内容均 Copy 自：https://mmdeploy.readthedocs.io/en/latest/

MMDeploy 提供了一系列工具，帮助用户更轻松的将 OpenMMLab 下的算法部署到各种设备与平台上。用户可以使用 MMDeploy 设计的流程一步到位，也可以定制自己的转换流程。

MMDeploy 定义的模型部署流程如下图所示：

主要分为以下几个部分：

• 模型转换（Model Converter）
  • 模型转换的主要功能是把输入的模型格式转换为目标设备的推理引擎所要求的模型格式。
  • 目前 MMDeploy 可以把 Pytorch 模型转换为 ONNX，TorchScript 等和设备无关的 IR 模型，也可以将 ONNX 模型转换为推理后端模型。两者相结合可实现端到端的模型转换也就是从训练端到生产端的一键式部署。
• MMDeploy 模型（MMDeploy Model）
  • 也称 SDK Model，它是模型转换结果的集合，不仅包括后端模型还包括模型的元信息，这些信息将用于推理 SDK 中
• 推理 SDK（Inference SDK）
  • 封装了模型的前处理、网络推理和后处理过程，对外提供多语言的模型推理接口。

1.3 RTMO

以下内容均 Copy 自：https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo

实时多人姿态估计在平衡速度和精度方面面临巨大挑战。两阶段自上而下的方法会随着图像中人数的增加而减慢速度，而现有的单阶段方法往往无法同时提供高精度和实时性。本文介绍的 RTMO 是一种单阶段姿态估计框架，它通过在 YOLO 架构中使用双一维热图表示关键点来无缝集成坐标分类，在保持高速的同时实现了与自上而下方法相当的精度。我们为热图学习提出了动态坐标分类器和量身定制的损失函数，专门用于解决坐标分类和密集预测模型之间不兼容的问题。RTMO 的性能优于最先进的单级姿态估计器，在 COCO 上的 AP 值提高了 1.1%，而在相同骨干网上的运行速度提高了约 9 倍。我们最大的模型 RTMO-l 在 COCO val2017 上达到了 74.8% 的 AP 值，在单个 V100 GPU 上达到了 141 FPS，证明了其效率和准确性。

更多细节大家可以查看论文：RTMO: Towards High-Performance One-Stage Real-Time Multi-Person Pose Estimation

2. 环境配置

在开始之前我们有必要配置下环境，mmpose 的环境可以通过 mmpose/Installation 文档中安装，mmdeploy 的环境也可以通过 mmdeploy/Installation 文档中安装

博主这里准备了一个可以运行 demo 和导出 ONNX 的环境，大家可以按照这个环境来，也可以自己参考文档进行相关环境配置

博主的环境安装指令如下所示：

conda create -n mmpose python=3.9
conda activate mmpose
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118
pip install -U openmim
mim install mmengine
mim install "mmcv>=2.0.0rc2"
mim install "mmpose>=1.1.0"
pip install mmdeploy==1.3.1
pip install mmdeploy-runtime==1.3.1

Note：这个环境博主目前只用于 demo 测试和 ONNX 导出，并不包含训练

为了不必要的错误，博主将虚拟环境中各个软件的版本都罗列出来，方便大家查看，环境如下：

Package                Version
---------------------- ------------
addict                 2.4.0       
aenum                  3.1.15      
aliyun-python-sdk-core 2.15.1      
aliyun-python-sdk-kms  2.16.3      
certifi                2022.12.7   
cffi                   1.16.0      
charset-normalizer     2.1.1       
chumpy                 0.70        
click                  8.1.7       
colorama               0.4.6       
coloredlogs            15.0.1      
contourpy              1.2.1       
crcmod                 1.7
cryptography           42.0.7      
cycler                 0.12.1      
Cython                 3.0.10      
dill                   0.3.8       
filelock               3.14.0      
flatbuffers            24.3.25     
fonttools              4.51.0      
grpcio                 1.64.0      
humanfriendly          10.0        
idna                   3.4
importlib_metadata     7.1.0       
importlib_resources    6.4.0       
Jinja2                 3.1.3       
jmespath               0.10.0      
json-tricks            3.17.3      
kiwisolver             1.4.5       
Markdown               3.6
markdown-it-py         3.0.0       
MarkupSafe             2.1.5       
matplotlib             3.9.0       
mdurl                  0.1.2
mmcv                   2.1.0
mmdeploy               1.3.1
mmdeploy-runtime       1.3.1
mmdet                  3.2.0
mmengine               0.10.4
mmpose                 1.3.1
model-index            0.1.11
mpmath                 1.3.0
multiprocess           0.70.16
munkres                1.1.4
netron                 7.6.8
networkx               3.2.1
numpy                  1.26.3
onnx                   1.16.0
onnx-simplifier        0.4.36
onnxruntime            1.18.0
opencv-python          4.9.0.80
opendatalab            0.0.10
openmim                0.3.9
openxlab               0.1.0
ordered-set            4.1.0
oss2                   2.17.0
packaging              24.0
pandas                 2.2.2
pillow                 10.2.0
pip                    24.0
platformdirs           4.2.2
prettytable            3.10.0
protobuf               3.20.2
pycocotools            2.0.7
pycparser              2.22
pycryptodome           3.20.0
Pygments               2.18.0
pyparsing              3.1.2
pyreadline3            3.4.1
python-dateutil        2.9.0.post0
pytz                   2023.4
pywin32                306
PyYAML                 6.0.1
regex                  2024.5.15
requests               2.28.2
rich                   13.4.2
scipy                  1.13.0
setuptools             60.2.0
shapely                2.0.4
six                    1.16.0
sympy                  1.12
tabulate               0.9.0
termcolor              2.4.0
terminaltables         3.1.10
tomli                  2.0.1
torch                  2.0.1+cu118
torchaudio             2.0.2+cu118
torchvision            0.15.2+cu118
tqdm                   4.65.2
typing_extensions      4.9.0
tzdata                 2024.1
urllib3                1.26.13
wcwidth                0.2.13
wheel                  0.43.0
xtcocotools            1.14.3
yapf                   0.40.2
zipp                   3.18.2

3. Demo测试

OK，环境准备好后我们就要开始执行 demo，具体流程可以参照：https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo

我们一个个来，首先是推理验证测试，教程给的推理脚本如下所示：

python demo/inferencer_demo.py $IMAGE --pose2d rtmo --vis-out-dir vis_results

在这之前我们需要把 mmpose 这个项目给 clone 下来，执行如下指令：

git clone https://github.com/open-mmlab/mmpose.git

也可手动点击下载，点击右上角的 Code 按键，将代码下载下来。至此整个项目就已经准备好了。

同时还要下载相关的预训练权重用于 Demo 测试和 ONNX 导出，RTMO 的预训练权重可以在 Model Zoo 中找到，如下所示：

可以看到 RTMO 的预训练权重非常多，博主这边选择的是红色框的这个 RTMO-s 的预训练权重进行后续的 Demo 测试和 ONNX 导出，点击后面的 ckpt 即可进行下载。

大家也可以点击 here 下载博主准备好的源码和权重（注意代码下载于 2024/6/1 日，若有改动请参考最新）

将下载好的预训练权重放在 mmpose 项目下，准备开始推理，执行如下指令即可进行推理：

python demo/inferencer_demo.py ./tests/data/coco/000000197388.jpg --pose2d ./configs/body_2d_keypoint/rtmo/body7/rtmo-s_8xb32-600e_body7-640x640.py --pose2d-weights ./rtmo-s_8xb32-600e_body7-640x640-dac2bf74_20231211.pth --vis-out-dir vis_results

Note：--pose2d 指定的配置文件一定要与 --pose2d-weights 相对应

输出如下所示：

执行成功后推理后的图片会保存在 mmpose/vis_results 文件夹下，推理结果如下图所示：

可以看到推理结果正常，效果还是 OK 的，下面我们就开始尝试将 RTMO 的 ONNX 导出来

4. ONNX导出初探

ONNX 的导出也可以参照：https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo

由于 ONNX 导出我们是在 mmdeploy 中完成的，因此需要先把 mmdeploy 这个项目给 clone 下来，执行如下指令：

git clone https://github.com/open-mmlab/mmdeploy.git

也可手动点击下载，点击右下角的 Code 按键，将代码下载下来。至此整个项目就已经准备好了。

同时还要下载相关的预训练权重用于 ONNX 导出，大家也可以点击 here 下载博主准备好的源码和权重（注意代码下载于 2024/6/1 日，若有改动请参考最新）

Note：我们需要将 mmpose 和 mmdeploy 放在同级目录下方便后续 ONNX 的导出，目录结构如下所示：

.
├─mmpose
├─mmdeploy

将下载好的预训练权重放在 mmdeploy 项目下，准备开始 ONNX 导出，执行如下指令即可进行导出：

python tools/deploy.py configs/mmpose/pose-detection_rtmo_onnxruntime_dynamic.py ../mmpose/configs/body_2d_keypoint/rtmo/body7/rtmo-s_8xb32-600e_body7-640x640.py ./rtmo-s_8xb32-600e_body7-640x640-dac2bf74_20231211.pth ../mmpose/tests/data/coco/000000197388.jpg --work-dir ./export

输出如下：

执行成功后导出的 ONNX 会保存在 mmdeploy/export 文件夹下，导出的 ONNX 如下图所示：

同时还保存着两张推理的图片，一张是 pytorch 推理的结果，一张是 onnx 推理的结果，如下图所示：

onnxruntime推理的结果

pytorch推理的结果

那大家可能会想这不就结束了吗？ONNX 已经导出来了呀，还有什么要做的吗？真有这么容易就好了，正因为导出的 ONNX 存在各种各样的问题，所以博主才写了这篇文章

那现在我们就一起来分析下导出的 ONNX 到底都存在着什么样的问题，我们用 Netron 查看这个 ONNX 会发现它前面还是挺正常的，后面简直是一团糟，博主这里截取了其中某些部分的网络结构，如下图所示：

那既然这样我们先用 onnx-simplifier 简化看看效果会不会好一点，新建简化脚本文件 onnx_sim.py，内容如下：

import onnx
import onnxsim

if __name__ == "__main__":

    model_onnx = onnx.load("./export/end2end.onnx")

    # 检查导入的onnx model
    onnx.checker.check_model(model_onnx)

    print(f"Simplifying with onnx-simplifier {onnxsim.__version__}...")
    model_onnx, check = onnxsim.simplify(model_onnx)
    assert check, "assert check failed"
    onnx.save(model_onnx, "./export/end2end.sim.onnx")

在终端执行下，指令如下：

python onnx_sim.py

我们再来看下简化后的 end2end.sim.onnx 有没有好点：

可以看到并没有什么作用，依旧是一团糟，这里我们就来看看为什么这个模型这么复杂，特别是后处理的部分，经过分析我们会发现如下的部分：

我们可以在 ONNX 模型中找到 NonMaxSuppression 这个算子，这显然就是把 NMS 等后处理部分一股脑全部塞到 ONNX 模型中了，所以模型看起来复杂度非常高，也难怪导出的 ONNX 名称是 end2end.onnx 即端到端的 ONNX 模型

不过这里博主吐槽一句，NonMaxSuppression 这种算子插到 ONNX 中干咩呀，也不加个 --end2end 的选项，好歹可以选是不是想要导出端到端的模型呀，真滴头疼😒

那我们再来看看还有哪些问题：

输入 input 维度没什么问题，输出的 dets 明显是边界框的坐标信息，keypoints 明显是边界框的关键点信息，为啥 keypoints 是四维的呢？把最后两个维度合并成一维不好吗？然后将 dets 和 keypoints 一起合并成一个输出不行吗？

OK，以上这些都是我们分析出来要去解决的，只有解决了这些问题我们才能够得到我们想要的 ONNX 模型，而不是这里导出的所谓的端到端的模型，不然的话你只能去使用 mmdeploy 这种东西去部署，耦合性太高了，也不怎么好用

因此总结下来我们现在的目的就是要解决两个问题：

• 找到模型 forward 中的 head 部分，剔除 NMS 部分
• 将输出节点 dets 和 keypoints 合并为一个输出

下面我们就一起来分析代码并解决掉这些问题

5. ONNX导出代码浅析

博主这里采用的是 vscode 进行代码的调试，其中的 launch.json 文件内容如下：

{
    // 使用 IntelliSense 了解相关属性。 
    // 悬停以查看现有属性的描述。
    // 欲了解更多信息，请访问: https://go.microsoft.com/fwlink/?linkid=830387
    "version": "0.2.0",
    "configurations": [
        {
            "name": "Python: 当前文件",
            "type": "python",
            "request": "launch",
            "program": "${file}",
            "cwd": "${workspaceFolder}",
            "console": "integratedTerminal",
            "args": [
                "configs/mmpose/pose-detection_rtmo_onnxruntime_dynamic.py",
                "../mmpose/configs/body_2d_keypoint/rtmo/body7/rtmo-s_8xb32-600e_body7-640x640.py", 
                "./rtmo-s_8xb32-600e_body7-640x640-dac2bf74_20231211.pth",
                "../mmpose/tests/data/coco/000000197388.jpg",
                "--work-dir", "./export",
            ],
            "justMyCode": true,
        }
    ]
}

要调试的文件是 mmdeploy 项目下的 tools/deploy.py，在 main 函数中打个断点我们来开始调试：

调试会发现我们最终会调用 torch2onnx 这个函数来进行 ONNX 的导出，因为这个函数执行完成后 ONNX 就已经导出来了，说明是在这个函数执行的过程中进行导出的

这个函数是在 mmdeploy/apis/pytorch2onnx.py 脚本文件中，但是博主尝试了调试根本跳不进去

没办法博主只好新建一个 export.py 的脚本文件来直接调用 torch2onnx 函数看里面到底执行了什么，内容如下：

import os
import argparse
from mmdeploy.apis.pytorch2onnx import torch2onnx

def parse_args():
    parser = argparse.ArgumentParser(description='Export model to backends.')
    parser.add_argument('deploy_cfg', help='deploy config path')
    parser.add_argument('model_cfg', help='model config path')
    parser.add_argument('checkpoint', help='model checkpoint path')
    parser.add_argument('img', help='image used to convert model model')
    parser.add_argument(
        '--test-img',
        default=None,
        type=str,
        nargs='+',
        help='image used to test model')
    parser.add_argument(
        '--work-dir',
        default=os.getcwd(),
        help='the dir to save logs and models')
    parser.add_argument(
        '--save-file',
        default=os.getcwd(),
        help='onnx model save name')    
    parser.add_argument(
        '--device', help='device used for conversion', default='cpu')

    args = parser.parse_args()
    return args

if __name__ == "__main__":
    args = parse_args()
    print(args)
    torch2onnx(args.img, args.work_dir, args.save_file, args.deploy_cfg, args.model_cfg, args.checkpoint, args.device)

launch.json 内容略有修改如下所示：

{
    // 使用 IntelliSense 了解相关属性。 
    // 悬停以查看现有属性的描述。
    // 欲了解更多信息，请访问: https://go.microsoft.com/fwlink/?linkid=830387
    "version": "0.2.0",
    "configurations": [
        {
            "name": "Python: 当前文件",
            "type": "python",
            "request": "launch",
            "program": "${file}",
            "cwd": "${workspaceFolder}",
            "console": "integratedTerminal",
            "args": [
                "configs/mmpose/pose-detection_rtmo_onnxruntime_dynamic.py",
                "../mmpose/configs/body_2d_keypoint/rtmo/body7/rtmo-s_8xb32-600e_body7-640x640.py",
                "./rtmo-s_8xb32-600e_body7-640x640-dac2bf74_20231211.pth",
                "../mmpose/tests/data/coco/000000197388.jpg",
                "--work-dir", "./export",
                "--save-file", "rtmo-s_8xb32-600e_body7-640x640.onnx"
            ],
            "justMyCode": true,
        }
    ]
}

我们来开始调试 export.py 脚本文件，调试后进入 torch2onnx 函数如下所示：

可以看到这个函数里面通过 build_pytorch_model 这个函数构建的 pytorch 的模型，那我们可以在这里就把 torch 模型导出来看看是不是我们想要的，新增代码如下所示：

torch_model = task_processor.build_pytorch_model(model_checkpoint)
# =====================================================================
import torch
dummy_input = torch.randn(1, 3, 640, 640)
torch.onnx.export(
    torch_model,
    dummy_input,
    "./export/model.onnx",
    export_params=True,
    input_names=["images"],
    output_names=["output0", "output1"],
    opset_version=11,
    dynamic_axes=None)    
# =====================================================================

我们一起来看看这个 model.onnx 模型，部分结构如下图所示：

可以看到模型非常非常简洁，就是我们想要的样子，但是问题是现在得到的输出是各个特征图的信息，还没有接入 head 部分变成我们真正想要的输出，所以我们只能接着往下走

我们可以看到真正调用的是 export 函数，它位于 mmdeploy/apis/onnx/export.py 脚本文件下，这其实就是一个套娃

在 export.py 中我们可以看到 model 经过了一个 patch_model 函数后变成了 patched_model，是不是这个函数里面执行了 head 的过程呢？我们可以把这个 patched_model 导出来看看长什么样子，新增如下代码：

patched_model = patch_model(model, cfg=deploy_cfg, backend=backend, ir=ir)
# =====================================================================
dummy_input = torch.randn(1, 3, 640, 640)
torch.onnx.export(
    patched_model,
    dummy_input,
    "./export/patched_model.onnx",
    export_params=True,
    input_names=input_names,
    output_names=output_names,
    opset_version=opset_version,
    dynamic_axes=None)
# =====================================================================

我们一起来看看这个 patched_model.onnx 模型，部分结构如下图所示

可以看到它和之前的 model 的结构一模一样，因此 head 部分的执行包括 NMS 算子的注册一定是在后面的代码中：

接着调试到最后我们发现了我们非常熟悉的 torch.onnx.export 函数，这个函数导出的模型就是 end2end.onnx，因此在 patch_model 和 torch.onnx.export 中间肯定是执行了一些操作，因此我们重点来分析下这中间的代码

代码如下所示：

if 'onnx_custom_passes' not in context_info:
    onnx_custom_passes = optimize_onnx if optimize else None
    context_info['onnx_custom_passes'] = onnx_custom_passes
with RewriterContext(**context_info), torch.no_grad():
    # patch input_metas
    if input_metas is not None:
        assert isinstance(
            input_metas, dict
        ), f'Expect input_metas type is dict, get {type(input_metas)}.'
        model_forward = patched_model.forward

        def wrap_forward(forward):

            def wrapper(*arg, **kwargs):
                return forward(*arg, **kwargs)

            return wrapper

        patched_model.forward = wrap_forward(patched_model.forward)
        patched_model.forward = partial(patched_model.forward,
                                        **input_metas)
    # force to export on cpu
    patched_model = patched_model.cpu()

博主是没咋看懂的，不过从函数命名来看应该是重写了 head 的 forward 部分进行了相关的重定位，博主让 ChatGPT 老师帮忙分析了一波，给出的解释如下：(from ChatGPT)

这段代码中的操作主要是围绕对 PyTorch 模型的前向传递（forward）方法进行修改，以便在导出模型到 ONNX 格式之前调整其行为。这里面包含几个关键步骤：

1. 使用 RewriterContext 和 torch.no_grad()：

• RewriterContext 可能是一个自定义的上下文管理器，用于在代码块内部调整或重写某些设置或状态。具体作用取决于 context_info 的内容，但通常这类上下文管理器用于临时修改环境或模型的配置。
• torch.no_grad() 是用于停用自动梯度计算的上下文管理器，这在模型推理或导出时常用，以减少内存消耗并提高计算速度，因为在这些场景下不需要计算梯度。

2. 条件检查并修改 input_metas：

• input_metas 如果不是 None，则断言其类型必须是字典。这意味着 input_metas 应包含一些元数据，这些元数据在模型前向传递中可能需要用到。
• 如果 input_metas 存在，代码将修改 patched_model.forward 函数，通过一个包装器 wrap_forward 来接管原有的 forward 方法。这个包装器实际上没有修改任何函数行为（只是简单地调用原始的 forward 并返回结果），但这种模式可以用来插入额外的逻辑或监控。

3. 修改 forward 方法以使用 input_metas：

• 经过 wrap_forward 包装后，再用 partial 函数固定一些关键字参数。partial 函数用于部分应用一个函数，固定原函数的一些参数值，返回一个新的可调用对象。这里通过 partial 将 input_metas 中的元数据作为默认参数值固定到 forward 方法上。

4. 将模型转移到 CPU：

• 通过调用 patched_model.cpu()，确保模型的所有参数和缓存都被转移到 CPU 内存中。这是在进行 ONNX 导出前的常见做法，以确保兼容性和避免任何与设备相关的问题。

综上所述，这段代码的主要目的是在导出前对模型的行为进行微调，确保输入的元数据能够被正确处理，并确保模型在 CPU 上进行导出。这些改动可能会影响导出的 ONNX 模型的结构或性能，尤其是如果 input_metas 包含对模型行为有重大影响的信息时。

经过博主的调试发现它最终会重定位到 mmdeploy/codebase/mmpose/models/heads/rtmo_head.py 的 predict 函数

注意这个函数才是我们真正需要修改的，它修改了 RTMOHead 类的 forward 函数的行为，执行了 head_module 以及 nms 等操作，下面我们先简单分析这里的代码大概都做了哪些事情

首先 Neck 出来的输入经过 self.module 得到预测结果，然后对结果展平，解码，如下图所示：

可以看到我们最终得到了五个预测输出：

• score：torch.Size([1, 2000, 1])
• flatten_bbox_preds：torch.Size([1, 2000, 4])
• flatten_pose_vecs：torch.Size([1, 2000, 256])
• flatten_kpt_vis：torch.Size([1, 2000, 17])
• bboxes：torch.Size([1, 2000, 4])

接着就是 NMS 部分如下图所示：

这部分就是把 NMS 部分插入到 ONNX 部分，然后根据 NMS 的结果过滤预测结果，其实这部分都没有必要放在 ONNX 里面，因为我们可以自己利用 CUDA 核函数实现

最后就是关键点解码，如下图所示：

可以看到模型最终有两个输出：

• dets：torch.Size([1, 11, 5])
  • 这个输出显然是框的信息
  • 其中 1 代表 batch 维度
  • 11 代表框的个数（经过 NMS 后）
  • 5 代表框的信息，注意是 left,top,right,bottom,conf 五个维度，而不再是我们熟悉的中心点宽高，这个我们可以从解码函数 bbox_xyxy2cs 中看出，这个是需要重点关注的，因为我们在自己实现后处理时会用到
• pred_kpts：torch.Size([1, 11, 17, 3])
  • 这个输出显然是关键点的信息
  • 其中 1 代表 batch 维度
  • 11 代表框的个数（经过 NMS 后）
  • 17 代表人体的 17 个关键点信息
  • 3 代表每个关键点的维度信息即 x,y,conf

OK，以上就是 RTMO 预测头 forward 重写的全部内容，下面我们就来看看该如何修改以满足我们的要求

6. 剔除NMS

首当其冲的肯定就是把 NMS 部分给干掉，我们找到对应的代码，并将其注释掉，如下所示：

# ========== rtmo_head.py ==========

# mmdeploy/codebase/mmpose/models/heads/rtmo_head.py第68行

# nms parameters
# post_params = get_post_processing_params(deploy_cfg)
# max_output_boxes_per_class = post_params.max_output_boxes_per_class
# iou_threshold = cfg.get('nms_thr', post_params.iou_threshold)
# score_threshold = cfg.get('score_thr', post_params.score_threshold)
# pre_top_k = post_params.get('pre_top_k', -1)
# keep_top_k = cfg.get('max_per_img', post_params.keep_top_k)

# do nms
# _, _, nms_indices = multiclass_nms(
#     bboxes,
#     scores,
#     max_output_boxes_per_class,
#     iou_threshold,
#     score_threshold,
#     pre_top_k=pre_top_k,
#     keep_top_k=keep_top_k,
#     output_index=True)

batch_inds = torch.arange(num_imgs, device=scores.device).view(-1, 1)

# filter predictions
dets = torch.cat([bboxes, scores], dim=2)
# dets = dets[batch_inds, nms_indices, ...]
# pose_vecs = flatten_pose_vecs[batch_inds, nms_indices, ...]
# kpt_vis = flatten_kpt_vis[batch_inds, nms_indices, ...]
# grids = self.flatten_priors[nms_indices, ...]
pose_vecs = flatten_pose_vecs
kpt_vis   = flatten_kpt_vis
grids     = self.flatten_priors

另外我们先导出静态 batch 的 ONNX 并加上 onnxsim 看看，我们始终坚持一个原则，那就是尽可能先将问题简单化，然后再去看别的，静态 ONNX 的修改如下：

# ========== export.py ==========

# mmdeploy/apis/onnx/export.py第149行

torch.onnx.export(
    patched_model,
    args,
    output_path,
    export_params=True,
    input_names=input_names,
    output_names=output_names,
    opset_version=opset_version,
    dynamic_axes=None,
    keep_initializers_as_inputs=keep_initializers_as_inputs,
    verbose=verbose)

# onnxsim 优化(新增)
# Checks
import onnx
model_onnx = onnx.load(output_path)
# onnx.checker.check_model(model_onnx)    # check onnx model

# Simplify
try:
    import onnxsim

    print(f"simplifying with onnxsim {onnxsim.__version__}...")
    model_onnx, check = onnxsim.simplify(model_onnx)
    assert check, "Simplified ONNX model could not be validated"
except Exception as e:
    print(f"simplifier failure: {e}")

onnx.save(model_onnx, output_path)
print(f"simplify done. onnx model save in {output_path}")

然后再次执行我们的导出脚本：

python export.py configs/mmpose/pose-detection_rtmo_onnxruntime_dynamic.py ../mmpose/configs/body_2d_keypoint/rtmo/body7/rtmo-s_8xb32-600e_body7-640x640.py ./rtmo-s_8xb32-600e_body7-640x640-dac2bf74_20231211.pth ../mmpose/tests/data/coco/000000197388.jpg --work-dir ./export --save-file rtmo-s_8xb32-600e_body7-640x640.onnx

输出如下图所示：

执行成功后会在 ./export 文件夹下生成 rtmo-s_8xb32-600e_body7-640x640.onnx，我们一起来看看这个 ONNX 的变化，部分结构如下图所示：

可以看到简洁了不少，NMS 被完全剔除了，这个就是我们想要的，不过还没有完全达到我们满意的程度，我们接着改

7. 输出合并

我们接着看，看下上面导出的 ONNX 的输入输出，如下图所示：

可以看到输出有两个，一个是 dets，一个是 keypoints，我们需要实现两个部分：

• keypoints 维度合并即 1x2000x17x3 变为 1x2000x51
• dets 和 keypoints 合并为一个输出即 1x2000x56

代码修改如下所示：

# ========== rtmo_head.py ==========

# mmdeploy/codebase/mmpose/models/heads/rtmo_head.py第98行

# decode keypoints
bbox_cs = torch.cat(bbox_xyxy2cs(dets[..., :4], self.bbox_padding), dim=-1)
keypoints = self.dcc.forward_test(pose_vecs, bbox_cs, grids)
pred_kpts = torch.cat([keypoints, kpt_vis.unsqueeze(-1)], dim=-1)
# 新增
bs, bboxes, ny, nx = pred_kpts.shape
pred_kpts = pred_kpts.view(bs, bboxes, ny*nx)

# return dets, pred_kpts
return torch.cat([dets, pred_kpts], dim=2)

同时输出合并成为一个了，导出时的节点名也顺便修改了，如下所示：

# ========== export.py ==========

# mmdeploy/apis/onnx/export.py第149行

torch.onnx.export(
    patched_model,
    args,
    output_path,
    export_params=True,
    input_names=["images"],  # 修改
    output_names=["output"], # 修改
    opset_version=opset_version,
    dynamic_axes=None,
    keep_initializers_as_inputs=keep_initializers_as_inputs,
    verbose=verbose)

再次执行下导出脚本看下导出的 ONNX，部分结构如下图所示：

可以看到输出合并成为一个了，符合我们的预期

8. LayerNormalization算子导出

还有可以优化的地方吗？有！我们来看上面导出的 ONNX 中某部分的结构：

大家有没有很熟悉呢？我们在韩君老师的课程中有讲过这个就是一个典型的 LayerNormalization 算子，大家感兴趣的可以看下：三. TensorRT基础入门-快速分析开源代码并导出onnx

那我们知道 ONNX 在 opset17 版本之后就开始支持 LayerNormalization 整个算子的导出了，具体可以参考：https://github.com/onnx/onnx/blob/main/docs/Operators.md

这里还有一个点需要大家注意，那就是 TensorRT 只有在 8.6 版本之后才开始支持 LayerNormalization 算子，因此如果你导出的 ONNX 中包含该算子，则需要你保证 TensorRT 在 8.6 版本以上，不然会出现算子节点无法解析的错误，具体可以参考：https://github.com/onnx/onnx-tensorrt/blob/release/8.6-EA/docs/Changelog.md

OK，那下面我们就将 opset 设置为 17 看看导出的 ONNX 有什么变化，代码修改如下：

# ========== export.py ==========

# mmdeploy/apis/onnx/export.py第149行

torch.onnx.export(
    patched_model,
    args,
    output_path,
    export_params=True,
    input_names=["images"],
    output_names=["output"],
    opset_version=17,
    dynamic_axes=None,
    keep_initializers_as_inputs=keep_initializers_as_inputs,
    verbose=verbose)

再来看看导出的 ONNX 又有什么样的变化，部分结构如下图所示：

可以看到之前的多个算子被完整的作为一个 LayerNormalization 算子导出来了，又优化了一步

9. 动态batch的实现

OK，上面导出的 ONNX 博主认为足够简洁已经没有什么问题了，既然这样我们就来看看动态 batch 下的 ONNX 导出是否会有一些不同

首先我们需要修改导出代码，让其保持 batch 维度的动态，如下所示：

# ========== export.py ==========

# mmdeploy/apis/onnx/export.py第149行

dynamic_batch = {'images': {0: 'batch'}, 'output': {0: 'batch'}}
torch.onnx.export(
    patched_model,
    args,
    output_path,
    export_params=True,
    input_names=["images"],
    output_names=["output"],
    opset_version=17,
    dynamic_axes=dynamic_batch, # 动态 batch
    keep_initializers_as_inputs=keep_initializers_as_inputs,
    verbose=verbose)

再次执行导出脚本看看导出的动态 ONNX 是否有新的变化，部分结构如下图所示：

可以看到导出的 ONNX 保持了 batch 维度的动态，其次整个 ONNX 没有什么多余的节点产生

这里多说一句，其实博主在另一个虚拟环境中（torch=1.12）导出动态 batch 的 ONNX 时会出现诸如 Shape、Gather 等节点存在，杜老师课程中讲过这个主要是 reshape 或者 view 是 batch 维度的值不是 -1 导致的，需要我们手动去修改代码保证 -1 在 batch 维度，断开跟踪，具体可以参考：6.3.tensorRT高级(1)-yolov5模型导出、编译到推理（无封装）

不过这里没有这个问题存在，大概率是 pytorch 在 2.0 版本之后做了一些优化更新使得导出的 ONNX 更加简洁

至此，我们完成了 RTMO 的 ONNX 导出，并将其修改成了我们期望的样子，下面我们来简单总结下都修改了哪些地方

10. 导出修改总结

经过上面的分析，我们来总结下要导出一个干净且符合我们要求的 RTMO 的 ONNX 到底需要做哪些修改

Note：以下操作都是在 https://github.com/open-mmlab/mmdeploy 项目中进行的

1. 修改 RTMO 预测头重写部分

# ========== rtmo_head.py ==========

# mmdeploy/codebase/mmpose/models/heads/rtmo_head.py第68行

# nms parameters
# post_params = get_post_processing_params(deploy_cfg)
# max_output_boxes_per_class = post_params.max_output_boxes_per_class
# iou_threshold = cfg.get('nms_thr', post_params.iou_threshold)
# score_threshold = cfg.get('score_thr', post_params.score_threshold)
# pre_top_k = post_params.get('pre_top_k', -1)
# keep_top_k = cfg.get('max_per_img', post_params.keep_top_k)

# do nms
# _, _, nms_indices = multiclass_nms(
#     bboxes,
#     scores,
#     max_output_boxes_per_class,
#     iou_threshold,
#     score_threshold,
#     pre_top_k=pre_top_k,
#     keep_top_k=keep_top_k,
#     output_index=True)

# batch_inds = torch.arange(num_imgs, device=scores.device).view(-1, 1)

# filter predictions
dets = torch.cat([bboxes, scores], dim=2)
# dets = dets[batch_inds, nms_indices, ...]
# pose_vecs = flatten_pose_vecs[batch_inds, nms_indices, ...]
# kpt_vis = flatten_kpt_vis[batch_inds, nms_indices, ...]
# grids = self.flatten_priors[nms_indices, ...]
pose_vecs = flatten_pose_vecs
kpt_vis   = flatten_kpt_vis
grids     = self.flatten_priors

# decode keypoints
bbox_cs = torch.cat(bbox_xyxy2cs(dets[..., :4], self.bbox_padding), dim=-1)
keypoints = self.dcc.forward_test(pose_vecs, bbox_cs, grids)
pred_kpts = torch.cat([keypoints, kpt_vis.unsqueeze(-1)], dim=-1)
bs, bboxes, ny, nx = pred_kpts.shape
pred_kpts = pred_kpts.view(bs, bboxes, ny*nx)

# return dets, pred_kpts
return torch.cat([dets, pred_kpts], dim=2)

2. 修改导出 ONNX 部分

# ========== export.py ==========

# mmdeploy/apis/onnx/export.py第138行

# torch.onnx.export(
#     patched_model,
#     args,
#     output_path,
#     export_params=True,
#     input_names=input_names,
#     output_names=output_names,
#     opset_version=opset_version,
#     dynamic_axes=dynamic_axes,
#     keep_initializers_as_inputs=keep_initializers_as_inputs,
#     verbose=verbose)

dynamic_batch = {'images': {0: 'batch'}, 'output': {0: 'batch'}}

torch.onnx.export(
    patched_model,
    args,
    output_path,
    input_names=["images"],
    output_names=["output"],
    opset_version=17,
    dynamic_axes=dynamic_batch
)

# Checks
import onnx
model_onnx = onnx.load(output_path)
# onnx.checker.check_model(model_onnx)    # check onnx model

# Simplify
try:
    import onnxsim

    print(f"simplifying with onnxsim {onnxsim.__version__}...")
    model_onnx, check = onnxsim.simplify(model_onnx)
    assert check, "Simplified ONNX model could not be validated"
except Exception as e:
    print(f"simplifier failure: {e}")

onnx.save(model_onnx, output_path)
print(f"simplify done. onnx model save in {output_path}")

3. 新建 export.py 导出文件

import os
import argparse
from mmdeploy.apis.pytorch2onnx import torch2onnx

def parse_args():
    parser = argparse.ArgumentParser(description='Export model to backends.')
    parser.add_argument('deploy_cfg', help='deploy config path')
    parser.add_argument('model_cfg', help='model config path')
    parser.add_argument('checkpoint', help='model checkpoint path')
    parser.add_argument('img', help='image used to convert model model')
    parser.add_argument(
        '--test-img',
        default=None,
        type=str,
        nargs='+',
        help='image used to test model')
    parser.add_argument(
        '--work-dir',
        default=os.getcwd(),
        help='the dir to save logs and models')
    parser.add_argument(
        '--save-file',
        default=os.getcwd(),
        help='onnx model save name')    
    parser.add_argument(
        '--device', help='device used for conversion', default='cpu')

    args = parser.parse_args()
    return args

if __name__ == "__main__":
    args = parse_args()
    print(args)
    torch2onnx(args.img, args.work_dir, args.save_file, args.deploy_cfg, args.model_cfg, args.checkpoint, args.device)

执行如下指令即可完成 RTMO 的 ONNX 导出

python export.py configs/mmpose/pose-detection_rtmo_onnxruntime_dynamic.py ../mmpose/configs/body_2d_keypoint/rtmo/body7/rtmo-s_8xb32-600e_body7-640x640.py ./rtmo-s_8xb32-600e_body7-640x640-dac2bf74_20231211.pth ../mmpose/tests/data/coco/000000197388.jpg --work-dir ./export --save-file rtmo-s_8xb32-600e_body7-640x640.onnx

导出好的 ONNX 文件在 export/rtmo-s_8xb32-600e_body7-640x640.onnx，这个也就是我们最终需要的 ONNX，大家也可以点击 here 下载博主导出好的 ONNX

注：前提是需要把环境配置好，模型权重文件下载好

11. 拓展-MMPose中导出ONNX

杜老师之前有教过如何在 MMDetection 中导出 YOLOX 的 ONNX 模型，大家可以感兴趣的可以看下：6.6.tensorRT高级(1)-mmdetection框架下yolox模型导出并推理

因此这次博主也尝试在 MMPose 中导出 RTMO 的 ONNX 模型，同时回顾下之前杜老师教的知识

开始之前我们需要将 mmpose 这个项目和对应的权重文件准备好

先验证下整个项目是否能成功，在 mmpose 项目下新建一个 export.py 的脚本文件，内容如下：

from mmpose.apis import init_model, inference_bottomup

img_path = 'tests/data/coco/000000000785.jpg'
config_file = "./configs/body_2d_keypoint/rtmo/body7/rtmo-s_8xb32-600e_body7-640x640.py"
checkpoint_file = "./rtmo-s_8xb32-600e_body7-640x640-dac2bf74_20231211.pth"

device = "cpu"

model = init_model(config_file, checkpoint_file, device=device)
inference_bottomup(model, img=img_path)

执行下该脚本，输出如下所示：

执行成功了，接下来我们就要去分析它，导出我们想要的 onnx，它的 model 是一个正常的 torch.model 的模型，因此我们直接导出看能不能成功，代码如下：

import torch
from mmpose.apis import init_model, inference_bottomup

img_path = 'tests/data/coco/000000000785.jpg'
config_file = "./configs/body_2d_keypoint/rtmo/body7/rtmo-s_8xb32-600e_body7-640x640.py"
checkpoint_file = "./rtmo-s_8xb32-600e_body7-640x640-dac2bf74_20231211.pth"

device = "cpu"

model = init_model(config_file, checkpoint_file, device=device)
# inference_bottomup(model, img=img_path)
torch.onnx.export(
    model,
    (torch.zeros(1, 3, 640, 640),),
    "model.onnx",
    opset_version=11
)

输出如下：

可以看到提示说 HybridEncoder 对象没有 pos_enc_0 属性，pos_enc_0 是什么呢？不清楚，很烦，这玩意没那么容易导出来，所以需要我们来进行分析

经过我们的调试分析（省略…😄）可以得知模型是需要 self.backbone、self.neck、self.head 这三项来完成推理的，所以我们完全可以自己来构建网络嘛，具体代码如下：

import torch
from mmpose.apis import init_model, inference_bottomup

img_path = 'tests/data/coco/000000000785.jpg'
config_file = "./configs/body_2d_keypoint/rtmo/body7/rtmo-s_8xb32-600e_body7-640x640.py"
checkpoint_file = "./rtmo-s_8xb32-600e_body7-640x640-dac2bf74_20231211.pth"

device = "cpu"

model = init_model(config_file, checkpoint_file, device=device)
x = torch.zeros(1, 3, 640, 640, device=device)
x = model.backbone(x)
x = model.neck(x)
x = model.head(x)
print(type(x))

# inference_bottomup(model, img=img_path)
# torch.onnx.export(
#     model,
#     (torch.zeros(1, 3, 640, 640),),
#     "model.onnx",
#     opset_version=11
# )

调试如下图所示：

大家对 x 的输出有没有很熟悉呢，它不就是我们之前没有进入预测头解码时导出的 ONNX 吗？之前导出的 ONNX 如下图所示：

那为什么这里导出来会提示说 HybridEncoder 对象没有 pos_enc_0 属性，那经过我们调试分析（省略…😄）其实在进行 neck 和 head 的 forward 之前都会去执行对应的 switch_to_deploy 函数，转换为部署模式提前计算出一些变量，比如这里的 pos_enc_0 就是 neck 的 Encoder 所需要的位置编码向量

既然如此我们就在 forward 之前来手动执行下 switch_to_deploy 函数，代码如下：

import torch
from mmpose.apis import init_model, inference_bottomup

img_path = 'tests/data/coco/000000000785.jpg'
config_file = "./configs/body_2d_keypoint/rtmo/body7/rtmo-s_8xb32-600e_body7-640x640.py"
checkpoint_file = "./rtmo-s_8xb32-600e_body7-640x640-dac2bf74_20231211.pth"

device = "cpu"

model = init_model(config_file, checkpoint_file, device=device)
class MyModel(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.model = init_model(config_file, checkpoint_file, device=device)
        test_cfg = {'input_size': (640, 640)}
        self.model.neck.switch_to_deploy(test_cfg)
        self.model.head.switch_to_deploy(test_cfg)
        self.model.head.dcc.switch_to_deploy(test_cfg)

    def forward(self, x):
        x = self.model.backbone(x)
        x = self.model.neck(x)
        x = self.model.head(x)

        return x

model = MyModel()
model.eval()

x = torch.zeros(1, 3, 640, 640, device=device)
torch.onnx.export(
    model,
    (x,),
    "model.onnx",
    opset_version=11,
    dynamic_axes=None
)

执行下该代码，输入如下：

导出的 ONNX 如下图所示：

导出的 ONNX 显然预测头的后处理没有做，因此我们需要将其梳理出来，整个后处理梳理过程有些繁琐，这边就不一一展示了，完整导出的代码如下所示：

import torch
from mmpose.apis import init_model
from mmpose.structures.bbox import bbox_xyxy2cs

class MyModel(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.model = init_model(config_file, checkpoint_file, device=device)
        test_cfg = {'input_size': (640, 640)}
        self.model.neck.switch_to_deploy(test_cfg)
        self.model.head.switch_to_deploy(test_cfg)
        self.model.head.dcc.switch_to_deploy(test_cfg)

    def forward(self, x):
        x = self.model.backbone(x)
        x = self.model.neck(x)
        cls_scores, bbox_preds, _, kpt_vis, pose_vecs = self.model.head(x)[:5]
        scores = self.model.head._flatten_predictions(cls_scores).sigmoid()
        flatten_bbox_preds = self.model.head._flatten_predictions(bbox_preds)
        flatten_pose_vecs  = self.model.head._flatten_predictions(pose_vecs)
        flatten_kpt_vis    = self.model.head._flatten_predictions(kpt_vis).sigmoid()
        bboxes = self.model.head.decode_bbox(flatten_bbox_preds, self.model.head.flatten_priors,
                                             self.model.head.flatten_stride)
        dets      = torch.cat([bboxes, scores], dim=2)
        grids     = self.model.head.flatten_priors
        bbox_cs   = torch.cat(bbox_xyxy2cs(dets[..., :4], self.model.head.bbox_padding), dim=-1)
        keypoints = self.model.head.dcc.forward_test(flatten_pose_vecs, bbox_cs, grids)
        pred_kpts = torch.cat([keypoints, flatten_kpt_vis.unsqueeze(-1)], dim=-1)
        bs, bboxes, ny, nx = map(int, pred_kpts.shape)
        bs = -1
        pred_kpts = pred_kpts.view(bs, bboxes, ny*nx)
        return torch.cat([dets, pred_kpts], dim=2)

if __name__ == "__main__":

    device = "cpu"
    config_file     = "configs/body_2d_keypoint/rtmo/body7/rtmo-s_8xb32-600e_body7-640x640.py"
    checkpoint_file = "rtmo-s_8xb32-600e_body7-640x640-dac2bf74_20231211.pth"

    model = MyModel()
    model.eval()

    x = torch.zeros(1, 3, 640, 640, device=device)
    dynamic_batch = {'images': {0: 'batch'}, 'output': {0: 'batch'}}
    torch.onnx.export(
        model,
        (x,),
        "rtmo-s_8xb32-600e_body7-640x640.onnx",
        input_names=["images"],
        output_names=["output"],
        opset_version=17,
        dynamic_axes=dynamic_batch
    )

    # Checks
    import onnx
    model_onnx = onnx.load("rtmo-s_8xb32-600e_body7-640x640.onnx")
    # onnx.checker.check_model(model_onnx)    # check onnx model

    # Simplify
    try:
        import onnxsim

        print(f"simplifying with onnxsim {onnxsim.__version__}...")
        model_onnx, check = onnxsim.simplify(model_onnx)
        assert check, "Simplified ONNX model could not be validated"
    except Exception as e:
        print(f"simplifier failure: {e}")

    onnx.save(model_onnx, "rtmo-s_8xb32-600e_body7-640x640.onnx")
    print(f"simplify done.")

导出的 ONNX 如下图所示：

可以看到和之前 mmdeploy 最终导出的 ONNX 如出一辙，大家也可以使用这种方法导出

在调试分析过程中博主对 RTMO 的结构有进一步的了解，如下所示：

• backbone：CSPDarknet
  • mmpose/models/backbones/csp_darknet.py
• neck：HybridEncoder 即 RT-DETR 的 Encoder 部分
  • mmpose/models/necks/hybrid_encoder.py
• head：RTMOHead
  • mmpose/models/heads/hybrid_heads/rtmo_head.py

结语

博主在这里对 mmpose 项目中的 RTMO 模型进行了 ONNX 导出，主要是利用 mmdeploy 导出，解决导出过程中出现的各种问题并达到我们最终想要的效果，此外博主还简单介绍了直接在 mmpose 下导出 ONNX，也能达到预期的效果。

总的来说 mmdeploy 和 mmpose 导出 ONNX 还是比较费劲的，封装得太死了，博主刚上手时曾一度想放弃，但是当你习惯之后问题总是可以解决的，不至于说束手无策

OK，以上就是 RTMO 的 ONNX 导出的全部内容了，下节我们来学习如何利用 tensorRT 推理 RTMO，敬请期待😄

下载链接

• 源代码、权重下载链接【提取码:rtmo】

参考

• RTMO: Towards High-Performance One-Stage Real-Time Multi-Person Pose Estimation
• https://github.com/open-mmlab/mmpose
• https://github.com/open-mmlab/mmdeploy
• https://mmpose.readthedocs.io/en/latest/
• https://mmdeploy.readthedocs.io/en/latest/
• https://github.com/open-mmlab/mmpose/tree/main/projects/rtmo
• mmpose/Installation
• mmdeploy/Installation
• 三. TensorRT基础入门-快速分析开源代码并导出onnx
• https://github.com/onnx/onnx/blob/main/docs/Operators.md
• https://github.com/onnx/onnx-tensorrt/blob/release/8.6-EA/docs/Changelog.md
• 6.3.tensorRT高级(1)-yolov5模型导出、编译到推理（无封装）
• 6.6.tensorRT高级(1)-mmdetection框架下yolox模型导出并推理