使用torchmetrics.detection的MeanAveragePrecision计算mAP遇长度不匹配错误

问题背景 #

基于PyTorch微调了fasterrcnn_resnet50_fpn模型做目标检测，使用torchmetrics.detection.MeanAveragePrecision计算验证集mAP时，运行代码抛出错误：

ValueError: Expected argument preds and target to have the same length, but got 100 and 127

相关代码片段：
推理函数：

@torch.no_grad
def generate_bboxes_on_one_img(image, model, device):
    model.to(device)
    model.eval()
    x = [image.to(device)]
    pred_boxes, pred_labels, pred_scores = model(x)[0].values()
    return pred_boxes, pred_labels, pred_scores

数据集与DataLoader设置：

def collate_fn(batch):
    return list(zip(*batch))

val_dataset = VisDroneDataset(val_images_path, val_annotations_df, transforms=val_transform)
val_data_loader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=0, collate_fn=collate_fn)

mAP计算代码：

mAP = MeanAveragePrecision(iou_type="bbox")
mAP.to(device)

for image, target in val_data_loader:
    original_boxes, original_labels, image_idx, _, _ = target[0].values()
    model.eval()
    x = [img.to(device) for img in image]
    preds_boxes, preds_labels, preds_scores = model(x)[0].values()

    image_PIL = val_dataset.get_image(image_idx)
    upscaled_image, pred_boxes_upscaled, labels = get_inverse_transform(image[0], 
                                                                    pred_boxes, 
                                                                    pred_labels, 
                                                                    *image_PIL.size)
    pred_to_mAP = [
            dict(
                boxes=torch.tensor(box, dtype=torch.float32),
                scores=score,
                labels=label
            )   for box, label, score in zip(pred_boxes_upscaled, pred_labels.clone().detach(), pred_scores.clone().detach())
            ]
            
    gt_to_mAP = [
            dict(
                boxes=original_boxes,
                labels=original_labels
            )   for box, label in zip(test_image_gt_bboxes, test_image_gt_labels)
            ]
    mAP.update(pred_to_mAP, gt_to_mAP)
    pprint(mAP.compute())

    break

一、MeanAveragePrecision工作原理 #

• 按批次图像进行指标累积：每次update传入的preds和target必须是长度等于批次图像数量的列表。
• 列表内每个元素对应单张图像的结果：比如批次size=1时，preds是[单图预测字典]，target是[单图标注字典]，两者长度均为1。
• 它会累积所有批次的预测与标注数据，最终调用compute时计算全局mAP，而非单张图像的mAP。

二、错误原因 #

1. 输入格式完全错误：你把单张图像的每个预测框/标注框都单独做成了字典，导致pred_to_mAP长度等于预测框数量（100），gt_to_mAP长度等于标注框数量（127），两者自然不匹配。
2. 误解输入结构：正确的格式是单张图像对应一个字典，字典内的boxes是该图像所有框的张量（形状(N,4)），labels是对应所有框的标签张量，scores是预测分数张量。

三、修正后的代码 #

mAP = MeanAveragePrecision(iou_type="bbox")
mAP.to(device)
# 提前把模型移到指定设备，避免重复操作
model.to(device)

for image, target in val_data_loader:
    # 提取当前图像的真实标注
    original_boxes, original_labels, image_idx, _, _ = target[0].values()
    original_boxes = original_boxes.to(device)
    original_labels = original_labels.to(device)

    model.eval()
    x = [img.to(device) for img in image]
    preds = model(x)[0]  # 获取单图的预测结果字典
    preds_boxes, preds_labels, preds_scores = preds.values()

    # 逆变换恢复预测框到原图尺寸
    image_PIL = val_dataset.get_image(image_idx)
    _, pred_boxes_upscaled, _ = get_inverse_transform(image[0], 
                                                      preds_boxes, 
                                                      preds_labels, 
                                                      *image_PIL.size)
    
    # 构造符合要求的preds输入：长度为1的列表，每个元素是单图的预测字典
    pred_to_mAP = [
        dict(
            boxes=torch.tensor(pred_boxes_upscaled, dtype=torch.float32).to(device),
            scores=preds_scores.clone().detach().to(device),
            labels=preds_labels.clone().detach().to(device)
        )
    ]
    
    # 构造符合要求的target输入：长度为1的列表，每个元素是单图的标注字典
    gt_to_mAP = [
        dict(
            boxes=original_boxes,
            labels=original_labels
        )
    ]
    
    # 更新指标
    mAP.update(pred_to_mAP, gt_to_mAP)

# 所有验证集图像跑完后，计算全局mAP
final_mAP = mAP.compute()
print(final_mAP)

额外注意事项 #

• 确保original_boxes与pred_boxes_upscaled的坐标格式一致（默认要求为xyxy格式），且均为(N,4)形状的张量。
• 不要在循环内重复移动模型到设备，提前执行一次即可。
• 建议在所有验证集批次都完成update后再调用compute，这样得到的是整个验证集的真实mAP。

内容的提问来源于stack exchange，提问作者ai.devmir