【mmdetection】解决Index put dtype不匹配：从Half到Float的NMS优化策略

1. 问题来了：高分辨率训练中的“数据类型不匹配”报错

最近在折腾mmdetection框架，想用Faster R-CNN配合Swin Transformer骨干网络训练一个目标检测模型。为了提高模型对小目标的识别能力，我决定把训练图片的尺寸调大一些。原本配置文件里train_pipeline中的Resize参数设置的是(1333, 800)，我心想，这分辨率是不是有点保守了？于是大手一挥，改成了(2666, 1600)，直接翻倍，想着这下模型能“看”得更清楚了。

结果，训练刚跑起来没多久，一个熟悉的“老朋友”——RuntimeError——就跳出来打招呼了。错误信息非常明确，但乍一看有点让人摸不着头脑：“Index put requires the source and destination dtypes match, got Half for the destination and Float for the source.” 翻译过来就是：索引赋值操作要求源张量和目标张量的数据类型必须匹配，但现在目标张量是Half（半精度浮点数，即FP16），而源张量是Float（单精度浮点数，即FP32）。这不匹配，所以程序罢工了。

更关键的是，错误堆栈清晰地指向了问题发生的具体位置：mmcv/ops/nms.py文件的第321行，具体是在batched_nms函数中，有一行代码试图将dets[:, -1]（推测是经过NMS后的得分）赋值给scores_after_nms[mask[keep]]。堆栈往上追溯，问题起源于RPN（区域提议网络）头部在进行边界框后处理时调用了这个NMS函数。简单来说，就是在模型推理（前向传播计算损失和预测结果）的过程中，进行非极大值抑制这一步时，数据类型对不上，卡壳了。

这个错误很有意思。它并不是在我一启动训练就立刻出现的，而是在训练过程中，当RPN网络生成了大量的候选框（proposals）后才触发的。这让我意识到，问题很可能与高分辨率图像输入有直接关系。图片尺寸变大，经过骨干网络和特征金字塔后，产生的特征图虽然可能经过了下采样，但相对于原图，其每个位置对应的感受野更精细，RPN网络在每个特征图位置上预设的锚框（anchors）数量是固定的，但特征图本身的分辨率（尺寸）可能因为输入变大而有所变化（取决于网络结构），或者更直接地说，高分辨率输入经过网络后，可能在某些阶段保留了更多的空间细节，导致RPN最终产生的初始提议框总数大大增加，从而引爆了这个隐藏在NMS计算中的数据类型“炸弹”。

2. 深入排查：错误根源与NMS的分批计算逻辑

看到错误堆栈指向mmcv中的NMS操作，我第一反应是去翻看源码。毕竟，知其然更要知其所以然，这样才能从根本上解决问题，而不是简单地回避。在mmcv/ops/nms.py文件中，我找到了batched_nms这个函数，并定位到了报错的大致区域。

为了让大家更清楚，我把相关的核心逻辑简化一下。在目标检测中，RPN或检测头会输出成千上万个预测框及其得分。NMS的作用就是过滤掉那些重叠度高且得分低的冗余框。当框的数量特别多时，为了效率和内存考虑，mmcv的实现里有一个分批处理的逻辑。它会根据每个框的类别ID（对于RPN阶段，通常所有框都属于“前景”这一类，但代码逻辑是通用的）将框分组，然后对每一组分别进行NMS。

问题就出在这个分批处理之后，对结果的赋值环节。我们来看一段模拟关键逻辑的代码：

# 假设 scores 是输入的所有框的得分，dtype 通常是 torch.float32 scores_after_nms = scores.new_zeros(scores.size()) # 创建一个和scores同形状、同dtype（Float）的全零张量 # 对每个唯一的类别id进行循环处理 for id in torch.unique(idxs): # 找出属于当前类别的框的索引 mask mask = (idxs == id) # 对这些框进行NMS计算，返回过滤后的框 dets 和保留的索引 keep dets, keep = nms_op(boxes_for_nms[mask], scores[mask], **nms_cfg_) # 将保留框的得分（dets[:, -1]）赋值给 scores_after_nms 中对应的位置 scores_after_nms[mask[keep]] = dets[:, -1] # ！！！报错发生在这里

报错信息明确指出，赋值操作左右两边的数据类型不匹配：scores_after_nms[mask[keep]]是目标，它是Half类型；而dets[:, -1]是源，它是Float类型。可是，我们明明看到scores_after_nms是用scores.new_zeros()创建的，理论上应该继承scores的数据类型（Float）。那么Half类型是从哪里冒出来的呢？

这就要提到现代深度学习训练中一个常用的加速技术：混合精度训练（Automatic Mixed Precision, AMP）。为了减少显存占用并加速计算，我们经常会使用AMP，它让模型的部分计算（尤其是卷积、矩阵乘等）在FP16（Half）精度下进行，同时保留一些关键操作（如损失计算、权重更新）在FP32（Float）精度下，以保持数值稳定性。

在mmdetection的配置文件中，开启混合精度训练通常是通过设置optim_wrapper.type = 'AmpOptimWrapper'来实现的。当AMP开启后，模型中的张量数据类型可能会在运行时动态变化。我推测，在错误发生的上下文中，scores_after_nms这个张量由于某种原因（可能是在之前的某个操作中被转换或创建时受到了AMP上下文的影响），其数据类型变成了Half。而dets[:, -1]这个从NMS操作内部返回的得分，可能由于NMS算子的实现或者输入数据的类型，仍然是Float。当程序试图将一个Float张量赋值给一个Half张量的切片时，PyTorch的严格类型检查就抛出了这个异常。

那么，为什么之前用(1333, 800)分辨率训练时没问题，换成(2666, 1600)就出问题了呢？关键可能在于框的数量。在batched_nms函数中，有一个参数控制着是否触发分批逻辑。当输入框的总数超过某个阈值（例如默认的10000）时，它才会进入上面的分批处理循环。在较低分辨率下，RPN产生的提议框数量可能低于这个阈值，直接进行全局NMS，赋值逻辑可能不同或者数据类型恰好一致。而高分辨率输入导致提议框数量激增，超过了阈值，从而走了分批处理的代码路径，恰好这条路径暴露了在混合精度环境下数据类型处理的不一致问题。

3. 解决方案一：调整训练配置（治标与权衡）

遇到问题，最直接的思路就是调整配置。这里有几个方向，各有优劣，你可以根据自己的项目需求和资源情况来选择。

第一个方案，关闭混合精度训练。这是最“粗暴”但最根本的解决方法。既然问题是Half和Float类型冲突引起的，那咱们不用Half精度不就行了？具体操作很简单，找到你的配置文件（通常是*.py文件），定位到optim_wrapper部分，将type从'AmpOptimWrapper'改为'OptimWrapper'。

# 修改前 optim_wrapper = dict(type='AmpOptimWrapper', ...) # 修改后 optim_wrapper = dict(type='OptimWrapper', ...)

这么一改，整个训练过程都会使用FP32精度，数据类型不一致的报错自然会消失。但是，代价也是明显的。FP32训练所需的显存大约是FP16的两倍，这意味着你的批量大小（batch size）可能不得不减小，或者同样的GPU能训练的模型尺寸变小了。更重要的是，计算速度会变慢，训练周期会显著拉长。对于大规模数据集和复杂模型，这个时间成本可能是难以接受的。所以，这个方案适合显存充足、对训练时间不敏感，或者只是想快速验证模型结构是否work的场景。

第二个方案，调整输入图像尺寸或NMS分批阈值。既然问题是高分辨率导致框数量过多触发了有问题的代码路径，那么我们可以从源头控制框的数量。

1. 调小输入尺寸：将train_pipeline中的Resize参数改回一个较小的值，比如(1333, 600)。这能立竿见影地减少RPN产生的锚框总数，使其低于分批阈值，从而绕过有问题的代码。缺点是对小目标检测不友好，违背了我们最初提升分辨率的目的。
2. 调大NMS分批阈值：在模型的训练配置中，找到RPN提议框的NMS设置。具体路径通常是model.train_cfg.rpn_proposal.nms。里面有一个参数叫max_num或者在某些版本中控制分批逻辑的阈值（可能需要查阅源码或具体版本的文档，有时是split_thr或通过max_num间接控制）。将这个值改得非常大（比如999999），使得几乎永远达不到分批条件，或者直接让框的数量永远低于阈值，从而避免进入那个有问题的赋值循环。缺点是如果框的数量真的非常多，单次进行NMS计算可能会消耗大量显存和计算资源，甚至可能引发OOM（内存溢出）。你需要权衡一下框的数量上限和你的硬件承受能力。

这些配置调整的方法，优点是不需要修改源代码，风险低。但缺点也很明显，它们是一种妥协和规避，要么牺牲了训练效率，要么牺牲了模型性能，要么增加了不稳定性。它们更像是临时救火，而不是从根本上修复漏洞。

4. 解决方案二：修改MMCV源码（治本与深入）

如果你追求一劳永逸，并且希望在高分辨率下稳定地使用混合精度训练，那么直接修改mmcv库的源码是最彻底的解决方案。这需要你对代码有一定的把控能力，并且记住，修改第三方库的源码可能会在将来库更新时被覆盖，需要重新处理或提交补丁。

我们的目标很明确：确保在batched_nms函数中，进行赋值操作时，源张量和目标张量的数据类型一致。根据错误位置，我们需要关注scores_after_nms[mask[keep]] = dets[:, -1]这一行。

思路一：在赋值前进行类型转换。这是最直接的修复方式。我们可以强制将dets[:, -1]转换为目标张量scores_after_nms的数据类型。

找到mmcv/ops/nms.py文件中的batched_nms函数（具体行号可能因版本不同略有差异，请根据错误堆栈定位），在报错行附近，添加一个类型转换操作：

# 修改前（大概在321行附近）： scores_after_nms[mask[keep]] = dets[:, -1] # 修改后： scores_after_nms[mask[keep]] = dets[:, -1].to(scores_after_nms.dtype)

这行代码的作用是，将dets[:, -1]这个张量转换（.to()）成与scores_after_nms相同的数据类型（.dtype），然后再进行赋值。这样无论scores_after_nms是Half还是Float，dets[:, -1]都会在赋值前与其对齐，从而避免类型不匹配的错误。

思路二：确保scores_after_nms的创建与上下文一致。另一种思路是审视scores_after_nms的创建方式。它是由scores.new_zeros()创建的。在混合精度训练中，如果scores张量在进入这个函数时已经被自动转换成了Half类型，那么scores_after_nms自然也是Half。问题可能在于NMS算子nms_op内部返回的dets其最后一列得分，并没有遵循同样的精度规则，或者在某些条件下保持了Float。在这种情况下，除了上述转换，我们也可以考虑是否应该在创建scores_after_nms时，就明确指定其数据类型与scores的原始数据类型或与NMS算子的预期输出类型保持一致？但这种方式侵入性更强，需要更仔细地分析数据流。

对于大多数情况，思路一的修改已经足够。它简单、有效，并且逻辑清晰：在数据融合点确保类型一致。修改后，保存文件，重新运行你的训练脚本。如果环境中的mmcv是通过pip install安装的，可能需要以“可编辑”模式重新安装，或者确保Python解释器加载的是你修改后的本地文件。

我个人的经验是，采用源码修改的方案后，之前在高分辨率下训练出现的Index putdtype错误再也没有出现过，混合精度训练得以顺利进行，既保证了训练速度，又实现了利用高分辨率提升小目标检测性能的初衷。当然，每次更新mmcv版本时，需要检查一下这个补丁是否仍然需要，或者是否已经被官方修复。

5. 实战演练与避坑指南

光说不练假把式，我们来模拟一个完整的实战场景，看看如何从零开始遇到并解决这个问题。

场景设定：你正在使用mmdetection v3.0.0 和 mmcv v2.1.0，在单张RTX 4090上训练一个基于Faster R-CNN with Swin-Tiny的自定义数据集检测模型。初始配置一切正常，输入尺寸为(1333, 800)，混合精度训练开启。为了提升对图像中细小文字的检测能力，你决定将输入尺寸提升至(2666, 1600)。

第一步：遭遇错误。修改配置文件中的train_pipeline的Resize步骤后，启动训练。几个iteration之后，程序崩溃，抛出前述的RuntimeError。错误堆栈清晰地指向mmcv/ops/nms.py。

第二步：分析判断。你迅速联想到这是高分辨率输入导致提议框数量增加，从而触发了batched_nms中特殊的分批处理路径。同时，注意到训练配置中optim_wrapper.type是'AmpOptimWrapper'，确认了混合精度训练是开启的。这基本坐实了“混合精度下数据类型不一致”的猜想。

第三步：选择方案。你评估了手头的选择：

• 方案A（关AMP）：显存占用会从18GB飙升至34GB，你的24GB显存扛不住，直接排除。
• 方案B（调小尺寸）：违背项目提升小目标检测的初衷，暂时不考虑。
• 方案C（调大阈值）：你尝试在配置文件中寻找model.train_cfg.rpn_proposal.nms，发现并没有直接的split_thr参数。查阅文档和源码发现，当前版本的batched_nms分批逻辑主要由max_num参数控制，但它的含义是最终保留的最大框数，并非严格的分批阈值。直接修改可能不奏效或行为不确定。
• 方案D（修改源码）：看起来是最直接、最符合需求的方案。

第四步：实施修改。

1. 找到你的Python环境下的mmcv安装位置。可以通过在Python中执行import mmcv; print(mmcv.__file__)来定位。
2. 打开对应的mmcv/ops/nms.py文件。
3. 使用搜索功能（Ctrl+F）查找scores_after_nms[mask[keep]] = dets[:, -1]。
4. 将其修改为scores_after_nms[mask[keep]] = dets[:, -1].to(scores_after_nms.dtype)。
5. 保存文件。

第五步：验证效果。重新启动训练脚本。观察日志，模型开始正常迭代，不再报错。使用nvidia-smi监控显存，发现依然保持在18GB左右，混合精度训练生效。训练速度与之前低分辨率时相差不大（因为计算量增大了，但AMP的加速效果仍在）。

避坑指南：

1. 版本差异：不同版本的mmcv和mmdetection，NMS的实现细节可能有差异。务必根据你的错误堆栈准确找到要修改的文件和行号，不要生搬硬套。
2. 环境隔离：如果你使用conda或venv管理多个项目环境，确保修改的是当前激活环境下的mmcv包。
3. 更新风险：当你未来通过pip install -U mmcv升级版本时，这个修改会被覆盖。建议保留你的修改记录，或者考虑向mmcv官方仓库提交一个Pull Request（如果确认是bug）。
4. 全面测试：修改后，不仅要在你的高分辨率训练任务上测试，如果可能，也跑一下原来的低分辨率配置，确保修改没有引入回归错误。
5. 备选方案：如果修改源码让你感到不安，或者是在共享的服务器环境中没有权限，可以退而求其次，尝试在配置文件中寻找是否还有其他隐藏参数可以控制NMS的行为，或者稍微降低一点目标分辨率（例如(2400, 1440)），使得框的数量刚好低于触发错误的分批阈值。

6. 理解背后：混合精度训练与类型传播

解决这个具体问题后，我们不妨深入一点，聊聊背后的原理。为什么混合精度训练会带来这些“烦人”的数据类型问题？理解这一点，能帮助我们在未来避免类似的坑。

混合精度训练的核心思想是：在保证训练精度（最终模型性能）不下降的前提下，尽可能使用FP16进行计算，以提升速度和节省显存。但是，FP16的数值表示范围（约 ±65504）和精度（约10^-4）远低于FP32（范围约 ±10^38，精度约10^-7）。这会导致两个主要问题：

1. 下溢：非常小的梯度（例如小于10^-5）在FP16中会变成0，导致权重无法更新。
2. 溢出：梯度或激活值太大，超过FP16的最大表示范围，变成无穷大（Inf）。

为了解决这些问题，AMP（如PyTorch的torch.cuda.amp）采用了一系列技术：

• 权重备份：保持一份FP32的模型权重主副本，用于更新。FP16的权重用于前向和反向传播。
• 损失缩放：将损失函数乘以一个缩放因子（如1024），等梯度放大后再反向传播，避免梯度下溢，然后在优化器更新权重前再缩放回来。
• 自动类型转换：AMP的autocast上下文管理器会自动将部分操作（如卷积、线性层）的输入转换为FP16，而将另一些操作（如softmax、损失函数）保持在FP32。

在mmdetection的AmpOptimWrapper中，就封装了这些逻辑。当开启后，模型的前向传播过程会在autocast上下文中进行。这意味着，张量的数据类型可能在操作之间动态变化。

现在，回到我们的NMS问题。scores张量作为输入进入batched_nms函数。它可能是在autocast区域外创建的（FP32），也可能在进入autocast区域后被转换成了FP16。而scores_after_nms = scores.new_zeros(scores.size())这一行，创建的新张量会继承scores的设备和数据类型。如果此时scores是FP16，那么scores_after_nms也是FP16。

关键在于nms_op（NMS算子）内部。这个算子可能是一个CUDA扩展，它的实现可能没有考虑到autocast上下文，或者其内部实现强制要求FP32计算以保证数值稳定性（比如排序、比较操作）。因此，无论输入是什么类型，它输出dets的最后一列得分（dets[:, -1]）可能固定为FP32。这就造成了赋值时的类型冲突。

所以，我们手动添加的.to(scores_after_nms.dtype)，实际上是在手动协调AMP自动类型转换机制与可能未充分适配AMP的第三方算子之间的数据类型鸿沟。这是一种非常实用的调试和修复思路：在自定义算子、复杂操作或框架边界处，留意数据类型的显式转换。

7. 扩展思考：其他可能的数据类型陷阱与排查方法

NMS中的这个Index put错误是一个典型案例，但在深度学习和mmdetection的使用中，数据类型不匹配的“坑”远不止这一个。了解常见的陷阱和排查方法，能让你在遇到类似问题时更加从容。

常见的数据类型陷阱：

1. 自定义算子的输入/输出类型：就像NMS算子一样，任何你自己编写的CUDA/C++扩展，或者直接使用某些底层PyTorch API实现的操作，如果没有仔细处理数据类型，在混合精度训练下都容易出问题。确保你的自定义算子能正确处理FP16输入，并输出一致的类型，或者明确文档说明其类型要求。
2. 模型初始化与加载：如果你用FP32精度预训练的权重去初始化一个准备用FP16训练的模型，要注意某些层（如BatchNorm的running_mean/var）可能需要保持FP32。加载检查点（checkpoint）时，如果保存时是混合精度状态，加载时环境不同也可能导致类型不匹配。
3. 损失函数与指标计算：一些复杂的损失函数或自定义的评价指标计算，可能涉及指数、对数、除法等对数值范围敏感的操作。在FP16下这些计算更容易产生Inf或NaN。通常这些部分会被AMP自动留在FP32上下文，但如果你的实现绕过了AMP，就需要手动处理。
4. 数据流水线（DataPipeline）：在mmdetection的train_pipeline中，Normalize、Pad等操作可能会改变数据的dtype。确保管道末端输入模型的数据类型是符合预期的（通常是FP32，但会被AMP自动转换）。

通用的排查与调试方法：

1. 打印张量信息：在怀疑出问题的地方，插入打印语句，输出关键张量的dtype、device和shape。这是最直接的调试手段。

   print(f"scores dtype: {scores.dtype}, scores_after_nms dtype: {scores_after_nms.dtype}") print(f"dets[:, -1] dtype: {dets[:, -1].dtype}")

2. 使用torch.autograd.detect_anomaly()：在训练脚本开始处启用异常检测，它能在发生NaN或Inf时提供更早的警告和堆栈信息，帮助你定位是哪个操作产生了数值问题。

   torch.autograd.set_detect_anomaly(True)

3. 梯度裁剪与损失缩放监控：在混合精度训练中，监控损失缩放因子（如果使用动态缩放）是否稳定。梯度爆炸有时也与类型问题相关。合理的梯度裁剪（clip_grad）可以增加稳定性。
4. 简化复现：如果问题复杂，尝试创建一个最小的、可复现的脚本。例如，单独提取出出错的NMS部分，用模拟的FP16和FP32数据喂给它，看是否报错。这能帮你快速确认问题核心。
5. 查阅框架文档与源码：最终，最权威的信息来源是官方文档和源码。了解mmcv中各个算子的设计，了解mmdetection中AMP包装器的具体行为，是避免和解决此类问题的根本。

记住，在深度学习工程中，数据类型、设备（CPU/GPU）、张量形状是三大基础要素，几乎所有的运行时错误都与之相关。养成在关键节点检查这些属性的习惯，能极大提升调试效率。这次解决NMS的dtype问题，正是这种思维的一次成功实践。当你下次再遇到类似的“神秘”报错时，不妨先问问自己：这里的张量，是什么类型？在哪里创建的？经过了哪些可能改变类型的操作？