集成Albumentations后MaskRCNN训练报KeyError: 'labels'求助

核心原因 #

PyTorch官方的MaskRCNN模型（包括maskrcnn_resnet50_fpn_v2）在训练时，要求每个target字典必须包含**labels**键，而非你自定义的class_labels。模型的ROI Heads模块会直接读取labels字段进行损失计算，找不到该键就会触发KeyError。

解决方案 #

1. 修改Dataset类中的target键名 #

在Dataset的__getitem__方法里，把target["class_labels"]替换为target["labels"]：

# 原代码
target["class_labels"] = class_labels

# 修改后
target["labels"] = class_labels

2. 同步Albumentations的label_fields参数 #

你当前的bbox_params中label_fields是空列表，需要将标签字段名加入，确保数据变换时标签能和bbox正确关联：

# 原代码
bbox_params = A.BboxParams(format='pascal_voc',label_fields = [])

# 修改后（与Dataset中标签变量名保持一致）
bbox_params = A.BboxParams(format='pascal_voc', label_fields=["class_labels"])

如果已经把target里的键改成labels，这里也可以直接写label_fields=["labels"]。

3. 额外优化点 #

• 训练循环中数据处理逻辑冗余，正确的处理方式应为：

for imgs, targets in train_loader:
    imgs = [img.to(device) for img in imgs]
    targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
    loss_dict = model(imgs, targets)
    # ...后续损失计算

• 你定义的loss_fn = nn.BCEWithLogitsLoss()并未实际使用，MaskRCNN会自动计算分类、框回归、mask等所有损失，无需额外定义损失函数。

修正后的关键代码片段 #

修改后的Dataset类__getitem__部分 #

## 创建类别标签
class_labels = torch.ones((num_objs,),dtype=torch.int64)

# Albumentations数据变换
if self.transform is not None:
  augmentations = self.transform(image=image, mask=mask, bboxes=boxes, class_labels=class_labels)
  image = augmentations["image"]
  mask = augmentations["mask"]
  boxes = augmentations["bboxes"]
  class_labels = augmentations["class_labels"]

boxes = torch.as_tensor(boxes, dtype = torch.float32)
masks = torch.as_tensor(masks, dtype=torch.uint8)

target = {}
target["boxes"] = boxes
target["masks"] = masks
# 修改键名为模型要求的labels
target["labels"] = class_labels

return image, target

修改后的bbox_params #

bbox_params = A.BboxParams(format='pascal_voc', label_fields=["class_labels"])

修改后的训练循环数据处理 #

for epoch in range(num_epochs):
    train_epoch_loss = 0
    val_epoch_loss = 0
    model.train()

    for imgs, targets in train_loader:
        imgs = [img.to(device) for img in imgs]
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
        loss_dict = model(imgs, targets)
        losses = sum(loss for loss in loss_dict.values())
        train_epoch_loss += losses.item()
        
        optimizer.zero_grad()
        losses.backward()
        optimizer.step()
    all_train_losses.append(train_epoch_loss)
    
    with torch.no_grad():
        model.eval()
        for imgs, targets in val_loader:
            imgs = [img.to(device) for img in imgs]
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
            loss_dict = model(imgs, targets)
            losses = sum(loss for loss in loss_dict.values())
            val_epoch_loss += losses.item()
        all_val_losses.append(val_epoch_loss)
    print(f"Epoch {epoch}: 训练损失 {train_epoch_loss}, 验证损失 {val_epoch_loss}")

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