PyTorch Faster R-CNN交通灯检测测试报训练模式需传targets错误

这个报错是torchvision内置Faster R-CNN模块的固定校验逻辑触发的：模块会通过自身的self.training布尔标识切换运行分支——

• 当self.training = True（训练模式）时，forward方法会强制要求传入真实标注的targets参数，用来计算分类、回归损失值
• 当self.training = False（推理模式）时，会跳过所有损失计算流程，仅接收图像输入，直接输出检测框、置信度、类别标签结果

当前的测试代码没有主动切换模型状态，因为之前训练循环里每轮都执行了model.train()，模型一直停留在训练态，所以测试时只传图像不传targets就会触发参数校验报错。问题和模型初始化、训练流程本身无关。

具体修正步骤 #

• 模型加载完成后，第一时间调用*model.eval()*把模式切换为推理态，从根源上关闭targets强制校验逻辑
• 推理阶段搭配*torch.no_grad()*关闭梯度计算，减少显存占用，避免推理过程意外修改模型参数
• 补全训练权重的加载逻辑，保证模型结构和训练时的结构完全一致（num_classes固定为4，不能随意改动）

修正后的可运行测试代码 #

# 按照训练时的配置初始化模型结构
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False)
N_CLASS = 4 
INP_FEATURES = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(INP_FEATURES, N_CLASS)

# 加载训练保存的权重文件
model.load_state_dict(torch.load('fasterrcnn_resnet{}_fpn.pth'.format(10)))
model.to(device)

# 核心修正：切换为推理模式
model.eval()
# 核心修正：推理阶段关闭梯度跟踪
with torch.no_grad():
    while(True):
        ret, input_frame = cap.read()
        if not ret:
            break
        image = input_frame.copy()
        input_tensor = preprocess(input_frame).float()
        input_tensor = input_tensor.unsqueeze_(0)
        input_tensor = input_tensor.type(torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor)

        # 推理阶段仅需传入图像，不需要额外传targets
        result = model(input_tensor)

        boxes = result[0]['boxes'].type(torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor)
        scores = result[0]['scores'].type(torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor)
        labels = result[0]['labels'].type(torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor)

        mask = nms(boxes, scores, 0.3)
        boxes = boxes[mask]
        scores = scores[mask]
        labels = labels[mask]

        boxes = boxes.data.cpu().numpy().astype(np.int32)
        scores = scores.data.cpu().numpy()
        labels = labels.data.cpu().numpy()

        mask = scores >= 0.5
        boxes = boxes[mask]
        scores = scores[mask]
        labels = labels[mask]

        for box, label in zip(boxes, labels):
            image = cv2.rectangle(image,
                              (box[0], box[1]),
                              (box[2], box[3]),
                              (0, 0, 255), 1)

        cv2.imshow("image", image)
        # 修正原逻辑中按任意键就退出的问题，改为按q键正常退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    # 退出后释放系统资源
    cv2.destroyAllWindows()
    cap.release()

容易踩的附加坑 #

1. Kaggle这类交互式Notebook环境会永久保留模型状态，训练完成后如果不手动调用model.eval()，模型会一直停留在训练态，后续所有推理操作都会触发相同的校验报错
2. 推理阶段不需要刻意构造空的targets参数传进模型，只要正确切换到eval模式，模型就不会对入参格式做训练态的校验
3. 加载权重时必须保证模型的最后预测层维度和训练时完全一致，否则会报权重维度不匹配的加载错误

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