使用Huggingface DETR训练COCO2017数据集时出现索引错误

错误根源 #

报错IndexError: index 64 is out of bounds for dimension 0 with size 3出现在DETR损失计算的class_cost = -out_prob[:, target_ids]行，核心原因是模型输出的类别维度与数据集的类别ID范围不匹配：

• COCO2017数据集包含90个实体类别+1个背景类别，类别ID范围为0~90；
• 但当前模型的分类头输出维度仅为3（继承自教程中针对balloon单类别数据集的配置），导致COCO数据中存在的类别ID（如64）超出了模型输出的索引范围。

解决步骤 #

1. 修正类别映射字典
   替换教程中针对balloon的id2label和label2id为COCO标准类别映射：

   # COCO完整类别映射可直接从官方定义获取，示例片段如下
   id2label = {
       0: 'background', 1: 'person', 2: 'bicycle', 3: 'car',
       4: 'motorcycle', 5: 'airplane', 6: 'bus', ..., 90: 'toothbrush'
   }
   label2id = {v: k for k, v in id2label.items()}
   

2. 重新初始化模型分类头
   加载预训练DETR模型时，指定匹配COCO的类别数量，覆盖教程中的单类别配置：

   from transformers import DetrForObjectDetection
   
   model = DetrForObjectDetection.from_pretrained(
       "facebook/detr-resnet50",
       num_labels=len(id2label),
       id2label=id2label,
       label2id=label2id
   )
   

   该操作会自动调整模型分类头的输出维度，使其与COCO的91个类别匹配。

3. 检查数据处理流程
   确保加载COCO数据集时，没有对原始类别ID进行错误重映射（比如教程中把balloon的类别ID统一改为1），保留COCO原生的ID规则（1~90对应实体类别，0对应背景）。

4. 验证维度一致性
   训练前可通过以下代码验证配置正确性：

   # 检查模型分类头输出维度
   print(f"模型输出类别数: {model.class_labels_classifier.out_features}")
   # 检查数据集中的最大类别ID
   sample_label = next(iter(train_dataloader))['labels'][0]['class_labels']
   print(f"数据集中最大类别ID: {sample_label.max().item()}")
   

   确保模型输出类别数（91）大于数据集中的最大类别ID（90）。

补充说明 #

此错误与是否使用GPU无关，完全是类别维度不匹配导致的索引越界。即使使用完整COCO数据集，只要模型分类头的类别数未适配，就会触发相同问题。

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