PyTorch目标检测训练阶段出现IndexError索引越界问题求助

问题根源 #

从你提供的批次数据可以明确看到，每个目标字典中的labels张量长度与boxes数量不匹配：

• 第一个target的boxes有5个边界框（shape为(5,4)），但labels仅为tensor([0])（shape为(1,)）
• 第二个target的boxes有3个边界框（shape为(3,4)），labels同样只有1个元素

这种维度不匹配会导致torchvision的ROI Heads在assign_targets_to_proposals方法中，尝试按边界框的索引去取对应标签时，出现索引越界（因为标签数量远少于边界框数量）。

解决步骤 #

• 检查数据集标注逻辑：确认你的标注文件中，每个边界框都对应一个标签。如果所有目标都属于同一类别（比如标签0），需要为每个边界框生成对应的标签，而不是只保留一个标签。
• 修复数据加载器的标签处理代码：在自定义数据集的__getitem__方法中，确保labels张量的长度与boxes的数量一致。例如，若所有box标签都是0，可通过labels = torch.tensor([0]*len(boxes), dtype=torch.int64)生成对应长度的标签张量。
• 验证修复后的数据：再次打印批次数据，确认每个target的labels第一维度大小与boxes的第一维度大小完全相同（比如5个box对应5个标签，3个box对应3个标签）。
• 额外排查点：如果使用了第三方标注工具转格式，检查转换脚本是否遗漏了部分标签的映射，导致仅保留了单个标签而非每个box对应一个标签。

内容的提问来源于stack exchange，提问作者user22267185