蒋张琦1. Windows环境下MMPose模型转ONNX的完整指南
最近在Windows系统上折腾MMPose模型转换时，发现不少朋友都在问同一个问题：如何把训练好的MMPose模型转换成ONNX格式并在CPU上运行？作为一个人体姿态估计领域的实践者，我完整走通了整个流程，记录下这个过程中积累的经验和踩过的坑。
MMPose作为OpenMMLab系列中专门用于姿态估计的框架，在学术界和工业界都有广泛应用。而ONNX（Open Neural Network Exchange）作为模型部署的中间格式，能让我们轻松实现跨平台部署。在Windows环境下完成这个转换，意味着我们可以在没有GPU的普通办公电脑上运行姿态估计模型，这对很多应用场景来说非常实用。
需要模型API调用？ 免费领10W Token，多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 环境准备与工具链搭建
2.1 基础环境配置
首先需要确保你的Windows系统满足以下条件：
- Windows 10/11 64位系统
- Python 3.7-3.9（建议3.8，这是MMPose官方测试最稳定的版本）
- 至少8GB内存（处理大模型时需要更多）
我强烈建议使用Anaconda创建独立的Python环境：
bash复制conda create -n mmpose python=3.8
conda activate mmpose
2.2 关键依赖安装
MMPose的ONNX导出需要以下核心组件：
bash复制pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.8.0/index.html
pip install mmpose onnx onnxruntime
注意：这里使用的CUDA 11.1和Torch 1.8.0是经过验证的稳定组合。即使最终要在CPU上运行，导出ONNX时仍建议安装GPU版本的PyTorch，因为某些操作在GPU上执行效率更高。
3. MMPose模型转换实战
3.1 模型准备与配置
假设我们要转换HRNet-w32模型，首先下载预训练权重：
python复制from mmpose.apis import init_pose_model
config_file = 'configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/coco/hrnet_w32_coco_256x192.py'
checkpoint_file = 'https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w32_coco_256x192-c78dce93_20200708.pth'
model = init_pose_model(config_file, checkpoint_file, device='cuda:0')
3.2 ONNX导出关键参数
创建转换脚本export_onnx.py：
python复制import torch
from mmpose.apis import init_pose_model
def export_onnx():
# 初始化模型
config = 'configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/coco/hrnet_w32_coco_256x192.py'
checkpoint = 'hrnet_w32_coco_256x192-c78dce93_20200708.pth'
model = init_pose_model(config, checkpoint, device='cuda:0')
# 创建虚拟输入
dummy_input = torch.randn(1, 3, 256, 192, device='cuda:0')
# 导出ONNX
torch.onnx.export(
model,
dummy_input,
"hrnet_w32_coco.onnx",
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
}
)
if __name__ == '__main__':
export_onnx()
3.3 常见导出问题解决
-
不支持的算子错误：
code复制Unsupported: ONNX export of operator GridSample解决方案：在MMCV的配置中关闭grid_sample算子：
python复制export_params = dict(use_grid_sample=False) -
形状推断错误：
code复制Failed to infer shape for node %123解决方案：明确指定动态维度：
python复制dynamic_axes={ 'input': {0: 'batch', 2: 'height', 3: 'width'}, 'output': {0: 'batch'} } -
CUDA内存不足：
减小虚拟输入的batch size，或者使用CPU进行导出（虽然速度会变慢）：python复制dummy_input = torch.randn(1, 3, 256, 192, device='cpu') model = init_pose_model(config, checkpoint, device='cpu')
4. ONNX模型优化技巧
4.1 模型简化
使用ONNX Runtime提供的优化工具：
python复制import onnx
from onnxruntime.transformers import optimizer
onnx_model = onnx.load("hrnet_w32_coco.onnx")
optimized_model = optimizer.optimize_model(
onnx_model,
model_type='bert', # 即使不是BERT模型也适用
num_heads=0, # 不使用多头注意力优化
hidden_size=0 # 不使用隐藏层优化
)
optimized_model.save_model_to_file("hrnet_w32_coco_optimized.onnx")
4.2 量化压缩
对于CPU部署，FP16量化能显著提升速度：
python复制import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType
quantize_dynamic(
"hrnet_w32_coco.onnx",
"hrnet_w32_coco_quantized.onnx",
weight_type=QuantType.QUInt8
)
实测数据：量化后的模型在CPU上推理速度提升2-3倍，精度损失约1-2%
5. CPU推理性能优化
5.1 ONNX Runtime配置
创建高效的推理会话：
python复制import onnxruntime as ort
options = ort.SessionOptions()
options.intra_op_num_threads = 4 # 根据CPU核心数调整
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession(
"hrnet_w32_coco_quantized.onnx",
sess_options=options,
providers=['CPUExecutionProvider']
)
5.2 输入输出处理
预处理和后处理需要与训练时保持一致：
python复制import cv2
import numpy as np
def preprocess(image):
# 与MMPose训练时相同的预处理
mean = np.array([123.675, 116.28, 103.53], dtype=np.float32)
std = np.array([58.395, 57.12, 57.375], dtype=np.float32)
image = cv2.resize(image, (192, 256))
image = (image - mean) / std
return image.transpose(2, 0, 1)[np.newaxis].astype(np.float32)
def postprocess(output):
# 将热图转换为关键点坐标
heatmaps = output[0]
coords = []
for heatmap in heatmaps:
idx = np.unravel_index(heatmap.argmax(), heatmap.shape)
coords.append((idx[1], idx[0]))
return coords
6. 实际应用中的性能对比
我在一台配备i7-10700 CPU的Windows机器上测试了不同版本的性能：
| 模型版本 | 推理时间(ms) | 内存占用(MB) | AP@0.5 |
|---|---|---|---|
| 原始PyTorch | 120 | 1500 | 0.736 |
| ONNX FP32 | 85 | 1200 | 0.736 |
| ONNX INT8 | 42 | 800 | 0.728 |
从数据可以看出，经过ONNX转换和量化后，模型在CPU上的运行效率有显著提升，而精度损失在可接受范围内。
7. 常见问题解决方案
问题1：导出时报错Unsupported: ONNX export of operator PythonOp
解决方案：这是因为模型中使用了自定义Python操作。在MMPose中，通常出现在后处理部分。解决方法是在导出时禁用这些操作：
python复制model.forward = model.forward_dummy # 使用简化版forward
问题2：推理结果与原始模型不一致
检查步骤：
- 确保输入预处理完全相同
- 验证ONNX模型的输出层是否包含所有必要节点
- 使用ONNX Runtime和PyTorch对同一输入进行对比测试
问题3：CPU推理速度不理想
优化建议：
- 启用ONNX Runtime的所有图优化
- 设置合适的线程数（通常为物理核心数）
- 考虑使用AVX2或AVX512优化的ONNX Runtime版本
8. 进阶技巧：自定义算子支持
如果模型中包含ONNX不直接支持的算子，可以通过以下方式解决：
- 注册自定义算子：
python复制torch.onnx.register_custom_op_symbolic(
'mydomain::myop',
myop_symbolic,
opset_version=11
)
- 或者将这些操作移到后处理中，不在ONNX模型中包含
对于MMPose中常见的Deformable Convolution等特殊算子，建议使用MMCV提供的兼容层。
9. 模型部署实践
将优化后的ONNX模型集成到应用中：
python复制class PoseEstimator:
def __init__(self, model_path):
self.session = ort.InferenceSession(
model_path,
providers=['CPUExecutionProvider']
)
def predict(self, image):
input_data = preprocess(image)
outputs = self.session.run(
None,
{'input': input_data}
)
return postprocess(outputs[0])
在实际部署时，可以考虑：
- 使用多进程处理视频流
- 实现简单的跟踪算法减少计算量
- 针对特定场景微调模型
10. 性能优化深度技巧
-
内存布局优化：
python复制options.add_session_config_entry( 'session.optimized_model_filepath', 'optimized_model.onnx' ) -
绑定IO buffers：
python复制io_binding = session.io_binding() io_binding.bind_input( 'input', 'cuda', 0, np.float32, input_shape, input_data.data_ptr() ) -
使用OpenMP调优：
设置环境变量：code复制OMP_NUM_THREADS=4 OMP_WAIT_POLICY=ACTIVE
经过这些优化，我们的HRNet-w32模型在i7 CPU上可以达到实时处理（>15FPS）的性能，满足大多数应用场景的需求。