mmdetection实例分割模型导出ONNX全流程避坑:动态shape与算子兼容实战
把训练好的实例分割模型导出成onnx本以为是临门一脚结果一脚踢在铁板上。这是我在实际部署mmdetection模型时遇到的最糟心问题之一模型在PyTorch里跑得好好的mAP也正常一旦走torch.onnx.export导出要么导出过程中直接报错要么导出成功但在onnx runtime里一推理就崩或输出全错。这个过程的坑密度极高而且很多报错信息不会直接告诉你“这里是动态shape问题”或者“这个算子没注册”需要一层层排查。这篇文章我不打算写一份官方文档式的教程而是把我实际排查过程中遇到的问题、定位思路、最终解法完整捋一遍。整个实例分割模型导出onnx无法运行的问题从头到尾大概可以拆成五个层面的坑导出链路本身、动态shape、实例分割特有的mask分支算子、NMS/RoIAlign这类黑盒算子的处理、以及导出成功之后的推理环境对齐。按这个顺序排查基本能覆盖绝大多数“onnx无法运行”的场景。1. 先看报错阶段Trace期、Export期、Runtime期完全不是一回事拿到“onnx无法运行”这个问题第一件事不是翻代码而是把报错发生的阶段搞清楚。不同阶段的报错指向的根因完全不同排查手段也天差地别。1.1 三个阶段的报错特征对比我习惯把问题分成三类导出时报错、导出成功但加载时报错、导出和加载都成功但推理结果不对。这三类问题的本质原因差别很大盲目在网上搜报错信息只会越搜越乱。报错阶段典型报错片段常见根因Trace期torch.onnx.export/ mmdeploy转换时TorchScript trace failed、Cannot insert a Tensor that requires grad模型forward里有动态控制流、训练状态未切换Export期onnx转换器执行时Unsupported operator、No symbolic registered某些自定义算子/高版本算子未注册onnx symbolicRuntime期onnx文件加载或推理时No Op registered、Input tensor shape mismatch、Got invalid input shapeopset版本太低、算子实现缺失、输入shape和固定shape不一致推理结果异常输出全零、mask乱成一团、检测框偏移预处理不一致、输出后处理遗漏、动态shape导致的reshape错位举例来说我在某次排查中碰到的情况是这样的导出命令执行到一半报了一大段类似RuntimeError: Exporting the operator grid_sampler to ONNX opset version 11 is not supported.的信息。当时第一反应是“换个opset版本”但换到13之后报错变成了另一条No Op registered for grid_sampler with domain_version of 13——这就说明问题已经从“转换器不支持”变成了“runtime端没有这个算子实现”。这种阶段判断能帮你少走两三个小时的弯路。1.2 一个高效的报错定位法二分裁剪法在动手改配置之前我强烈建议先做一个“最小导出验证”把模型固定到单张输入、固定尺寸暂时屏蔽掉不必要的分支只保留主干结构再导出。如果最小结构能导出再逐步加回mask分支、加回动态尺寸。我遇到过很多项目组的小伙伴一出错就翻github issue其实问题就出在某个不常用分支的某个算子上二分裁剪法很快就能锁死。具体做法也很简单用mmdetection官方配置里的configs加载模型不改任何参数只把model.train_cfg置空确保模型处于eval模式。构造一个固定shape的dummy input比如torch.randn(1, 3, 800, 800)。先用torch.jit.trace走一遍确认TorchScript层面能不能过。再过torch.onnx.export把报错去掉的部分一层层加回来。这个流程跑一遍之后你对“到底哪一层出的问题”就会非常清楚。下面说的所有坑都是在我做这种排查时一个个撞出来的。2. 头号杀手动态shape导致导出崩溃或推理错乱大多数mmdetection模型在训练时都是动态shape的图片可以任意尺寸batch size也可以随意设置。但onnx导出这件事本质上是对一个静态计算图的序列化。如果你在导出时没有把输入shape固定下来或者代码里有依赖输入shape的reshape/插值逻辑导出器就会面临一个两难要么把动态逻辑硬编码成静态逻辑要么报错。2.1 固定shape是导出实例分割模型的默认前提我见过不少新手的导出脚本dummy input直接用了训练时的数据加载器逻辑图片尺寸不固定。这会导致几种情况torch.onnx.export会警告“The shape inference of ONNX is not fully supported”但最终还是导出了。导出的onnx里某些层的shape变成了[1, 256, -1, -1]即动态维度。在onnx runtime里一旦输入尺寸和导出的固定尺寸不一致报Invalid input shape或直接silent fail。这里要明白一件事onnx的dynamic_axes参数虽然可以让模型接受不同尺寸输入但代价是几乎所有算子都得按动态shape走很多针对静态shape的图优化都做不了。而且对实例分割模型来说mask分支的resize插值特别敏感动态shape和插值算子叠加在一起很容易导出成一个“看起来能跑跑起来全错”的文件。我在实践中建议部署场景如果不需要多种分辨率输入就直接固定输入尺寸。比如固定为(1, 3, 800, 800)或(1, 3, 1280, 800)一劳永逸。如果一定要支持多个尺寸那也别用dynamic_axes传一堆维度而是固定宽高比、只允许batch维动态或者干脆导多个onnx文件一个尺寸一个文件。后一种方案在工业界非常常见简单可靠。2.2 一个典型的动态shape导出行错误reshape和数据错位举例来说mask分支在经过RoIAlign之后得到的是一个[num_rois, C, mask_size, mask_size]的特征图后续要经过几次卷积和反卷积最后再插值到原图尺寸。如果num_rois是动态的那模型内部就会有大量依赖num_rois的reshape和view操作。在onnx导出过程中最怕遇到的就是view和reshape的shape参数是变量。因为onnx的Reshape算子的输出形状要求要么是个常量要么是计算得到的tensor当这个tensor的值在trace时依赖于某个非const的中间结果就很容易导出失败或者导出后推理时因为实际输入和trace时不一致reshape把数据搞错位。我记得有一次导出的onnx在onnx runtime推理时屡屡出现“Detection mask尺寸完全乱掉”的现象。排查到最后发现是某个分支里用了x.view(x.size(0), -1)这样的写法x.size(0)在trace时被绑定到了dummy input的num_rois数量比如100但真实推理时num_rois变成了300于是reshape后的特征图就会错位。这个问题的解法倒不复杂要么在导出时把RPN的top_k参数调大确保推理时的候选框数量不超过trace时的数量要么干脆把view改成flattenonnx对flatten的shape推断更robust一些。但要注意这不是改个API名字那么简单你需要理解模型内部在哪里用了动态维度。3. 实例分割特有雷区Mask分支的grid_sampler和插值算子如果你跑的是Cascade Mask R-CNN或Mask R-CNN这类两阶段实例分割模型导出失败的重灾区几乎都集中在mask分支。这部分涉及的算子主要有两个一个是RoIAlign一个是mask上采样时用的grid_sampler/ interpolate。检测分支和分类分支就算能导出mask分支一加上去就崩这是实例分割导出的经典剧本。3.1 Mask分支为什么要“特殊对待”实例分割和纯目标检测的核心差异就是多了一个mask分支它为每个RoI预测一个二值掩码。标准做法是从RoIAlign得到固定尺寸的特征后通过几层卷积得到每个类别的mask logits然后通过torch.nn.functional.interpolate或F.grid_sample把这个logits采样到原图对应的位置。问题出在grid_sampler这个算子上。PyTorch里的F.grid_sample映射到ONNX时在opset 11及以下没有官方支持在opset 13以上才部分支持。而mmdetection某些版本的实现里mask分支恰好会走这个逻辑。你可能会问为什么不能用interpolate因为mask要在RoI坐标系里做对齐不是简单地把整个feature map放大而grid_sample这种按坐标采样的方式更适合这种场景。实测中我遇到的问题是这样的RuntimeError: Exporting the operator grid_sampler to ONNX opset version 11 is not supported.换到opset 13又变成onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph: [ONNXRuntimeError] No Op registered for grid_sampler with domain_version of 13这就是典型的“导出器支持了但runtime端没实现”的情况。那一刻真的有点崩溃因为报错信息本身没有任何“怎么解决”的提示。3.2 处理grid_sampler不兼容的三种方案针对这个坑我试过三种方案各有取舍。第一种在导出时把F.grid_sample换成F.interpolate。但问题是这两者的语义并不完全一致换完之后mask精度可能下降而且如果模型已经在训练阶段用了grid_sample推理时突然换interpolate结果不会等同。第二种使用mmdeploy的导出工具链。mmdeploy针对mmdetection的mask分支做了不少算子补偿其中就包括对grid_sampler的处理。实测下来它会把grid_sampler转换成一个等效的子图组合从而绕开算子缺失的问题。但代价是onnx文件变得更复杂推理速度会略受影响而且在某些老版本onnx runtime上依然可能出问题。第三种也是最robust的方案自定义onnx symbolic把grid_sampler映射到一个组合算子序列。具体做法是在导出时注册一个symbolic函数用aten::slice、aten::cat、aten::mul等基础算子来模拟grid_sample的前向计算。这个方法代码量不大但对onnx的算子语义要够熟不然很容易写出一个“能导出但结果全错”的映射。我实际采用的就是第三种自己写了一个简化版的grid_sampler_symbolic因为我们的mask分辨率固定为28x28RoI的数量在导出时也做了固定所以整个映射写起来并不复杂。如果你也要走这条路我的建议是先检查你用的PyTorch和onnxruntime版本再查一下grid_sampler在你那个opset下到底有没有support如果确定没有再决定要不要自己写。版本越新这个问题越少但老项目里是真的会遇到的。4. 检测头的黑盒算子NMS和RoIAlign的注册与替换除了mask分支实例分割模型的检测头部分也有两个“黑盒算子”很容易让导出失败。一个是NMS非极大值抑制一个是RoIAlign。前者是后处理逻辑后者是特征提取的核心算子。这两个算子在PyTorch里实现都很快但在onnx导出时特别麻烦。4.1 NMS到底该留在模型里还是留在模型外先说说NMS。标准的Faster R-CNN / Mask R-CNN推理流程中模型内部有一个BBoxHead里面会做类别分类和bbox回归之后再通过NMS过滤重复框。如果你直接用原model导出NMS会作为一个算子出现在计算图中。但PyTorch的NMS在onnx里并没有一个一对一的官方映射——torchvision.ops.nms导不出标准ONNX算子。mmdetection的nms操作往往是自己封装的底层可能是torchvision.ops.nms也可能是自定义C实现。这类操作在trace时往往被“内联”成一个大的Trace节点导出成onnx后onnx runtime根本不认。我碰到过好几次这样的情况导出时没有报错但onnx文件打开一看里面有一个孤零零的ATen算子或是一大段无法解析的Scope节点。这就是NMS没有被正确映射的后果。对于这个问题我建议不要试图把NMS放进onnx模型里。实例分割的NMS逻辑远比单纯的目标检测复杂——每个类别要单独NMS、mask的IoU也需要参与判断稍有不慎就容易出差错。更务实的做法是导出时只保留模型的前向推理部分也就是直到输出boxes、scores、labels和mask logits为止的部分把NMS作为后处理放在onnx runtime之外。这样导出的onnx文件会在最后多输出一些原始候选框后处理再用Python或C代码做NMS。从工程上讲这个方案的稳定性和可调试性都远高于硬把NMS塞进图里。如果你确实需要端到端的onnx导出那可以考虑在导出前把nms换成ONNX标准里的NonMaxSuppression算子。但这种做法的前提是你的模型结构恰好能套上这个算子输入格式对于mmdetection这类有自定义multiclass_nms逻辑的模型来说很少能直接套用需要大量的兼容改造性价比不高。4.2 RoIAlign的自定义算子注册实战再来看RoIAlign。这个算子在onnx的标准算子列表里其实有对应项即RoiAlign但它对输入格式有严格要求——坐标必须是[batch_idx, x1, y1, x2, y2]的形式。而mmdetection内部处理的RoI坐标往往经过了自己的一套变形逻辑直接用标准RoiAlign算子替代时容易出现格式对不上的问题。我当时处理这个问题的路径是这样的先打印出模型里RoIAlign的输入输出shape和数值分布确认我拿到的RoI坐标和标准RoiAlign期望的输入格式是否一致。如果一致那么在导出时注册一个自定义symbolic把torchvision.ops.roi_align映射到ONNX的RoiAlign算子。如果不一致就先在导出前用一个lambda包装层把坐标转换好再调用标准算子。这里要特别强调一点RoIAlign的参数spatial_scale、sampling_ratio、aligned必须和训练时保持完全一致。很多导出后推理结果不准的案例问题不是算子没导出成功而是aligned这个参数缺了——它在onnx标准算子里的默认值和torchvision的实现不一致会导致mask和box的输出偏差零点几个像素看起来不明显但在实例分割这种像素级任务上影响会被mask插值放大。这里贴一段我当时注册自定义symbolic的参考代码片段不是完整代码只是示意思路from torch.onnx import register_custom_op_symbolic def roi_align_symbolic(g, input, rois, spatial_scale, output_size, sampling_ratio, aligned): # 将torchvision的roi_align映射到onnx的RoiAlign算子 return g.op(RoiAlign, input, rois, output_height_ioutput_size, output_width_ioutput_size, sampling_ratio_isampling_ratio, spatial_scale_fspatial_scale, coordinate_transformation_mode_shalf_pixel)有了这个symbolic之后导出时torchvision.ops.roi_align就能顺利映射了。这里的output_size需要是int或者tuplesampling_ratio要小心如果训练时设的是0表示自适应导出时就要显式设成一个正整数不然onnx端的默认行为会和PyTorch不一致。5. 导出成功也不代表完事ONNX Runtime推理环境对齐如果说前四步是“从0到1”的导出过程那么第五步就是“从1到100”的推理验证。我见过太多人倒在导出这一步也见过更多人倒在导出成功之后——onnx加载成功、跑起来了但输出结果和PyTorch完全对不上或者时对时不对。这种问题最恶心因为没有任何报错只有一个“不对”排查全靠猜。5.1 输入预处理必须一字不差地复刻在onnxruntime里跑推理输入数据预处理这一步看似简单实际非常容易出问题。mmdetection的测试流程里包含一整套图像处理LoadImageFromFile、Resize、Pad、Normalize、DefaultFormatBundle。其中有一个非常隐蔽的坑test pipeline里的Pad策略。mmdetection在测试时默认会pad到divisor32的整数倍也就是让图像宽高对齐到32的倍数。你如果用原图比如800x600直接喂给onnx模型而模型导出时的dummy input是固定(1, 3, 800, 800)的话那么输入尺寸不匹配还只是第一层问题。更隐蔽的是如果你用动态shape导出模型内部的padding值可能已经根据导出时dummy input的形状写死了。实际的解决方法是导出前分析test pipeline把预处理流程完整复刻一遍。我当时就写了一个预处理函数包含cv2.resize或PIL.Image.resize把短边resize到800同时保持长宽比。按照divisor32做padding。用model.data_preprocessor里的mean和std做归一化。你可以直接从mmdetection的test_pipeline里把Normalize的mean/std抄出来但要注意归一化必须在padding之后还是应该在padding之前这个顺序在mmdetection里是由pipeline定义的不是你想当然的。错了的话模型推理结果会差得非常离谱因为padding区域的像素值会被当成正常图像内容参与特征提取。5.2 输出解析的几个隐蔽深坑实例分割模型导出的输出一般包含检测框、类别、分数和mask logits。在onnxruntime里拿到这些输出之后还有一个后处理的过程。这一步有三处容易出错第一mask logits需要经过sigmoid再resize到原图尺寸。很多人忘记sigmoid直接把logits当概率用导致mask看起来像被水泡过一样。第二mask的resize要根据检测框的位置裁剪。mmdetection的mask_head输出的是每个RoI对应的mask logits尺寸是[num_rois, num_classes, mask_size, mask_size]你需要根据检测框在原图中的坐标把这个mask对应区域裁剪出来再放大到原图大小。如果直接整张图resize输出就会在空间位置上完全错位。第三box坐标系的偏移。mmdetection在测试时会使用get_bboxes里面包含一个clip操作和一个offset修正。这个offset和我们之前说的padding有关——如果padding了那么输出的坐标是相对于padding后的图的你必须减掉padding的尺寸才能得到原图坐标。否则框的位置会整体偏右下看起来像是“对的”但细看全错了。我自己的做法是先用一张小图、一个单对象的简单场景做对比验证。在PyTorch里推理一次把boxes、scores、labels、masks全部打印出来然后在onnxruntime里用完全相同的输入推理把同样结果打印出来。逐项对比看差异出现在哪一步。这种“黄金对比法”不需要什么花哨工具就是最笨也最有效的方法。6. 如果还是搞不定用mmdeploy替代手搓导出说实话mmdetection从训练到onnx导出这一整套事情官方是提供了一套相对完整的工具链的也就是mmdeploy。如果你不想手写symbolic也不想逐个处理NMS、RoIAlign、grid_sampler的兼容问题直接用mmdeploy可能是投入产出比最高的方案。6.1 mmdeploy能把哪些坑帮你填平mmdeploy把前面的几个坑都做了处理它会自动检测模型结构将mmdetection里的MulticlassNMS替换成ONNX可导出的组合逻辑或者将其标记为后处理。它对mask分支的grid_sampler和RoIAlign做了算子等价替换和覆盖。它提供了onnxruntime、tensorrt等多个后端同一个模型可以导出多种runtime格式。我实际测试过一个原本在torch.onnx.export下怎么都导不出来的Mask R-CNN在mmdeploy的onnxruntime后端下可以相对顺利地完成转换。当然mmdeploy也有自己的“脾气”比如它对mmdetection版本和PyTorch版本有匹配要求版本一乱转出来的文件照样是坏的。6.2 mmdeploy的坑不是装上就能用使用mmdeploy时最常见的错误就是版本不匹配。mmdetection从2.x到3.xconfig和模型结构变动不少mmdeploy的适配也只覆盖特定版本组合。如果你拿新版的mmdetection模型搭配一个老版本mmdeploy导出时可能直接报KeyError或ModuleNotFoundError。我踩过的另一个坑是mmdeploy导出时对torch.jit.trace的依赖很重而实例分割模型的mask分支里如果有动态控制流比如根据RoI数量改变循环次数trace过程就会失败。这时mmdeploy会提示你用--dynamic参数但实际动态模式下onnx推理速度会明显下降而且某些算子仍可能不兼容。所以我的看法是mmdeploy是一个很好的工具适合“模型结构比较常规、版本匹配没问题”的场景。如果你已经确定自己的模型结构比较特殊比如加了自定义head、改了loss和postprocess那还是老老实实走手动导出加后处理外置的路线虽然费时但可控性强出了问题也容易追溯。7. 一些防止再踩坑的通用建议最后分享几条我总结出来的经验算是对整个排查过程做一个收束。第一导出前先做“模型简化”。把推理时用不到的hook、可视化和后处理模块全部摘掉只保留model.backbone、model.neck和model.roi_head这几个核心组件。可以直接在mmdetection的model配置上做裁剪导出的计算图越小出问题的面就越少。第二导出时固定opset版本不要盲目追最新。onnxruntime的老版本不一定支持新opset导出的文件新版本又可能改变某些算子的默认行为。我的建议是使用opset 13或14这两个版本覆盖了大多数常用算子的稳定性区间。不要用opset 17、18这些太新的版本除非你确定你的runtime完全支持。第三onnx文件导出来之后至少用两种runtime各跑一遍。我一般在导出后先用onnxruntime的CPU版本做基本验证再用TensorRT或openvino做一遍推理对比。这两种runtime在算子的实现上存在细节差异如果两边结果一致那模型本身基本没问题如果不一致就说明某个算子在这个runtime上没有被完全支持或实现有偏差需要提前处理否则部署上去就是大雷。第四所有后处理代码必须写单元测试。这一点容易被很多人忽视特别是在时间紧的时候。由于onnx导出的输出往往不是最终的boxes和masks需要不少后处理操作而每个操作都可能引入细微偏差。把这些后处理封装成函数输入固定的onnx输出tensor断言期望的box坐标和mask尺寸能让你在换模型、换尺寸、换runtime时快速发现回归。我自己的习惯是任何部署项目里都会保留一套和模型训练时完全一致的预处理和后处理代码用同一张测试图在PyTorch和onnxruntime上做对比。这张测试图的选择也有讲究——不要用一张刚训练集的图片而是用一张包含多尺度目标、密集场景的真实图片。因为单目标简单图很可能掩盖掉动态shape、padding这些隐藏问题。多目标、多尺度的复杂图才能让问题暴露出来。导出实例分割模型的onnx文件确实比单纯的目标检测模型复杂不少。但吃透上面这几层问题之后你会发现“无法运行”这个宽泛的描述其实只是少数几个核心矛盾的不同表现形式动态shape与静态图的矛盾、黑盒算子和onnx标准算子的矛盾、预处理与后处理的环境对齐矛盾。把这三组矛盾解决掉剩下的就是版本匹配和参数细节的打磨了。