PyTorch FasterRCNN预训练模型无有效检测输出问题求助

模型返回无效预测是两个核心错误叠加导致的：

• 随机初始化的检测头覆盖了预训练权重：加载官方预训练权重后，你直接新建了随机初始化的FastRCNNPredictor替换原模型的box_predictor，把预训练好的分类、边界框回归头全部丢弃。随机初始化的检测头没有经过训练，无法输出正确预测。你定义的CLASSES列表和COCO官方91类（含背景、预留无效位）完全一致，不需要替换检测头。
• 预处理逻辑不符合预训练数据分布：torchvision的Faster RCNN预训练模型，输入要求是像素值范围在[0,1]的浮点型张量，不需要额外做ImageNet均值方差归一化，也不需要额外做min-max缩放。你先做ImageNet归一化把像素值拉到0均值附近，再强行min-max缩放回0-1区间，完全打乱了像素分布，和模型预训练时的输入分布不匹配，必然输出错误结果。另外read_image读入的默认是uint8类型张量，像素值范围为[0,255]，你直接转成FloatTensor后没有除以255，初始数值范围本身就不对。

1. 模型加载阶段，删掉替换box_predictor的代码。如果后续要训练自定义类别，再替换检测头，且替换后必须加载自己训练好的权重才能推理，不能直接用官方预训练权重跑推理。
2. 推理阶段的预处理简化为：将uint8类型的输入张量转成浮点型，除以255把像素值缩放到0-1区间即可，去掉多余的归一化、min-max缩放操作。
3. 补全visualizer模块缺失的导入，避免运行报错。

detector.py #

import torch
import torchvision
from config import device
from torchvision.models.detection import FasterRCNN_ResNet50_FPN_V2_Weights
from torchvision.models.detection.faster_rcnn import fasterrcnn_resnet50_fpn_v2

CLASSES = ['__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

class detector():
  def __init__(self, threshold=0.3, model_path=None):
    self.threshold = threshold
    self.model = self.load_model(model_path)
  def load_model(self, model_path=None):
    print("Loading model...")
    model = fasterrcnn_resnet50_fpn_v2(weights=FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT)
    print("Loading model...done")

    # 自定义类别训练时再取消下面注释，且训练完成后才能用于推理
    # in_features = model.roi_heads.box_predictor.cls_score.in_features
    # num_classes = len(CLASSES)
    # model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) 

    if model_path is not None:
      print("Loading model from:", model_path)
      model.load_state_dict(torch.load(model_path, map_location=device))

    model.eval()
    return model

  def detect(self, image):
    """
    image: read_image读入的C*H*W格式uint8张量
    """
    # 预处理仅做类型转换和像素值缩放，无额外操作
    image_tensor = image.type(torch.FloatTensor) / 255.0

    # 单张图片补充batch维度
    if len(image_tensor.shape) == 3:
      image_tensor = image_tensor.unsqueeze(0)

    image_tensor = image_tensor.to(device)
    self.model.to(device)

    with torch.no_grad():
      outputs = self.model(image_tensor)

    outputs = [{k: v.to('cpu') for k, v in t.items()} for t in outputs]

    boxes = outputs[0]['boxes']
    scores = outputs[0]['scores']
    labels = outputs[0]['labels']

    # 按置信度阈值过滤结果
    conf_mask = scores > self.threshold
    boxes = boxes[conf_mask]
    labels = labels[conf_mask]

    pred_classes = [CLASSES[i] for i in labels.cpu()]

    return boxes, pred_classes

visualizer.py #

from torchvision.io import read_image, write_png
from torchvision.utils import draw_bounding_boxes
from torch.utils.tensorboard import SummaryWriter

def image_visualizer(detector,img_path, tensorboard=False):
    print('Loading image...')
    img_tensor = read_image(img_path) # C*H*W格式uint8张量
    boxes, pred_cls = detector.detect(img_tensor)
    if boxes.shape[0] == 0:
        print("No objects detected")
        return

    bbox_img = draw_bounding_boxes(img_tensor, boxes, labels=pred_cls, width=2)

    if tensorboard:
        print("Saving output to tensorboard...")
        writer = SummaryWriter()
        writer.add_image('detection_result', bbox_img)
        writer.close()
        print("Image saved to tensorboard")
    else:
        write_png(bbox_img, 'detection_img.png')
        print("Detection result saved to detection_img.png")

调用示例 #

from detector import detector
from visualizer import image_visualizer

image_path = "sailboat.jpg"
predictor = detector(threshold = .3, model_path = None)
image_visualizer(predictor, image_path, tensorboard=False)

之前测试“不修改box_predictor层也返回错误结果”，完全是预处理逻辑错误导致的，和检测头无关。torchvision的目标检测预训练模型，预训练阶段没有使用ImageNet均值方差归一化，不要把分类模型的预处理逻辑套到检测模型上。

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