Torchvision RetinaNet 预测不需要的班级背景

Question

我想用我的自定义数据集和 2 个类（无背景）训练来自 torchvision 的预训练 RetinaNet。为了使用 RetinaNet 进行训练，我做了以下修改：

num_classes = 3 # num of objects to identify + background class

model = torchvision.models.detection.retinanet_resnet50_fpn(pretrained=True)
# replace classification layer 
in_features = model.head.classification_head.conv[0].in_channels
num_anchors = model.head.classification_head.num_anchors
model.head.classification_head.num_classes = num_classes

cls_logits = torch.nn.Conv2d(in_features, num_anchors * num_classes, kernel_size = 3, stride=1, padding=1)
torch.nn.init.normal_(cls_logits.weight, std=0.01)  # as per pytorch code
torch.nn.init.constant_(cls_logits.bias, -math.log((1 - 0.01) / 0.01))  # as per pytorcch code 
# assign cls head to model
model.head.classification_head.cls_logits = cls_logits

问题是我得到了 0 类的检测，无论 num_classes 是 2 还是 3，它都是背景。

我尝试理解源代码并找不到类似的内容 fasterrcnn roi_head

# remove predictions with the background label
boxes = boxes[:, 1:]
scores = scores[:, 1:]
labels = labels[:, 1:]

我该如何解决这个问题？任何帮助将不胜感激！