在RetinaNet中微调PyTorch模型

Question

为了创建我自己的对象检测，我想完善火炬视觉中可用的经过预先训练的RetinaNet模型。

我试图复制在这个链接上为FastRCNN所做的工作：https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html#finetuning-from-a-pretrained-model

我所做的是：

model = model = torchvision.models.detection.retinanet_resnet50_fpn(pretrained=True)
num_classes = 2 

# get number of input features and anchor boxed for the classifier
in_features = model.head.classification_head.conv[0].in_channels
num_anchors = model.head.classification_head.num_anchors

# replace the pre-trained head with a new one
model.head = RetinaNetHead(in_features, num_anchors, num_classes)

模型被宣布了，训练也没有中断。然而，表现如此糟糕，以至于一个非常愚蠢的检测都不起作用。

我的问题是，我编写的代码可以重新训练RetinaNet模型吗？

Answer 1

我也试图做一件类似的事情。下面的代码应该有效。在COCO数据集上加载预先训练的权重后，我们需要用自己的分类器层替换分类器层。

num_classes = # 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(out_channels, 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

回归盒网络不需要更改，因为每个空间位置的锚盒数量与模型在COCO数据集上预培训时一样没有变化。与更快的R相比，代码有点麻烦.希望看到一个更优雅的解决方案。

希望这能有所帮助。