小样本目标检测（FSOD）之FADI代码解析

相对于最原始的faster rcnn做了如下7处改动
1.使用wordnet计算与新类最相似的类别
2.在roi_head.bbox_head上新增加了CosineSimBBoxHead类（来自TFA的cos头）
3.冻结某些层unfreeze_layers
4.新增加FewShotVOCDataset和FewShotVOCTestDataset的dataset类
-----------------------discrimination---------------
5.修改了rpn_head为FADIRPNHead，适应不同的框架
6.在roi_head.bbox_head中新增了FADIBBoxHead
7.新增加了SetSpecializedMarginLoss损失

改动二：CosineSimBBoxHead类

import torch
import torch.nn as nn
from mmdet.models.builder import HEADS
from mmdet.models.roi_heads.bbox_heads import ConvFCBBoxHead


@HEADS.register_module()
class CosineSimBBoxHead(ConvFCBBoxHead):
    def __init__(self,
                 fc_out_channels=1024,
                 scale=20.,
                 with_margin=False,
                 *args,
                 **kwargs):
        super(CosineSimBBoxHead,
              self).__init__(num_shared_convs=0,
                             num_shared_fcs=2,
                             num_cls_convs=0,
                             num_cls_fcs=0,
                             num_reg_convs=0,
                             num_reg_fcs=0,
                             fc_out_channels=fc_out_channels,
                             *args,
                             **kwargs)
        self.fc_cls = nn.Linear(self.cls_last_dim,
                                self.num_classes + 1,
                                bias=False)
        self.scale = scale
        self.with_margin = with_margin

    def forward(self, x, return_fc_feat=False):
        x = x.flatten(1)
        for fc in self.shared_fcs:
            x = self.relu(fc(x))

        # normalize the input x along the `input_size` dimension
        x_norm = torch.norm(x, p=2, dim=1).unsqueeze(1).expand_as(x)
        x_normalized = x.div(x_norm + 1e-5)

        # normalize weight
        temp_norm = torch.norm(self.fc_cls.weight.data, p=2,
                               dim=1).unsqueeze(1).expand_as(
                                   self.fc_cls.weight.data)
        self.fc_cls.weight.data = self.fc_cls.weight.data.div(temp_norm + 1e-5)
        cos_dist = self.fc_cls(x_normalized)
        scores = self.scale * cos_dist
        bbox_preds = self.fc_reg(x)
        if return_fc_feat:
            return scores, bbox_preds, x_normalized
        return scores, bbox_preds

    def forward_cls(self, x):
        x = x.flatten(1)

        for fc in self.shared_fcs:
            x = self.relu(fc(x))

        # normalize the input x along the `input_size` dimension
        x_norm = torch.norm(x, p=2, dim=1).unsqueeze(1).expand_as(x)
        x_normalized = x.div(x_norm + 1e-5)

        # normalize weight
        temp_norm = torch.norm(self.fc_cls.weight.data, p=2,
                               dim=1).unsqueeze(1).expand_as(
                                   self.fc_cls.weight.data)
        self.fc_cls.weight.data = self.fc_cls.weight.data.div(temp_norm + 1e-5)
        cos_dist = self.fc_cls(x_normalized)
        scores = self.scale * cos_dist
        return scores

    def forward_bbox(self, x):
        x = x.flatten(1)

        for fc in self.shared_fcs:
            x = self.relu(fc(x))

        bbox_preds = self.fc_reg(x)
        return bbox_preds

    def init_weights(self):
        # conv layers are already initialized by ConvModule
        if self.with_cls:
            nn.init.normal_(self.fc_cls.weight, 0, 0.01)
        if self.with_reg:
            nn.init.normal_(self.fc_reg.weight, 0, 0.001)
            nn.init.constant_(self.fc_reg.bias, 0)