〖MMDetection〗解析文件：mmdet/models/dense_heads/anchor_free_head.py

《解析 OpenMMLab 中 Anchor-Free Head 的实现》

在计算机视觉目标检测领域，Anchor-Free 的检测方法近年来受到了广泛关注。OpenMMLab 中的mmdet/models/dense_heads/anchor_free_head.py文件实现了一种通用的 Anchor-Free 检测头，为各种具体的 Anchor-Free 模型（如 FCOS、Fovea、RepPoints 等）提供了一个基础框架。

一、文件概述

这个文件定义了AnchorFreeHead类，它继承自BaseDenseHead。主要功能是在不依赖预定义锚框的情况下进行目标检测，通过对输入特征图进行处理，预测目标的类别和边界框。

二、主要类和方法解析

1. AnchorFreeHead类的初始化

@MODELS.register_module()
class AnchorFreeHead(BaseDenseHead):
    def __init__(
        self,
        num_classes,
        in_channels,
        feat_channels=256,
        stacked_convs=4,
        strides=(4, 8, 16, 32, 64),
        dcn_on_last_conv=False,
        conv_bias='auto',
        loss_cls=dict(
            type='FocalLoss',
            use_sigmoid=True,
            gamma=2.0,
            alpha=0.25,
            loss_weight=1.0),
        loss_bbox=dict(type='IoULoss', loss_weight=1.0),
        bbox_coder=dict(type='DistancePointBBoxCoder'),
        conv_cfg=None,
        norm_cfg=None,
        train_cfg=None,
        test_cfg=None,
        init_cfg=dict(
            type='Normal',
            layer='Conv2d',
            std=0.01,
            override=dict(
                type='Normal', name='conv_cls', std=0.01, bias_prob=0.01))
    ):
        super().__init__(init_cfg=init_cfg)
        # 设置各种参数，如类别数量、输入通道数、特征通道数等
        self.num_classes = num_classes
        self.use_sigmoid_cls = loss_cls.get('use_sigmoid', False)
        if self.use_sigmoid_cls:
            self.cls_out_channels = num_classes
        else:
            self.cls_out_channels = num_classes + 1
        self.in_channels = in_channels
        self.feat_channels = feat_channels
        self.stacked_convs = stacked_convs
        self.strides = strides
        self.dcn_on_last_conv = dcn_on_last_conv
        assert conv_bias == 'auto' or isinstance(conv_bias, bool)
        self.conv_bias = conv_bias
        self.loss_cls = MODELS.build(loss_cls)
        self.loss_bbox = MODELS.build(loss_bbox)
        self.bbox_coder = TASK_UTILS.build(bbox_coder)
        # 创建点生成器，用于确定特征图上的采样点
        self.prior_generator = MlvlPointGenerator(strides)
        self.num_base_priors = self.prior_generator.num_base_priors[0]
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
        self.conv_cfg = conv_cfg
        self.norm_cfg = norm_cfg
        self.fp16_enabled = False
        # 初始化各层
        self._init_layers()

在初始化函数中，设置了模型的各种参数，包括类别数量、输入通道数、损失函数、框编码器等。同时，通过MlvlPointGenerator创建了点生成器，用于后续确定特征图上的采样点。最后调用_init_layers方法初始化模型的各个层。

2. _init_layers方法

    def _init_layers(self):
        self._init_cls_convs()
        self._init_reg_convs()
        self._init_predictor()

这个方法分别调用了三个内部方法来初始化分类卷积层、回归卷积层和预测器层。

2.1 _init_cls_convs方法

    def _init_cls_convs(self):
        self.cls_convs = nn.ModuleList()
        for i in range(self.stacked_convs):
            chn = self.in_channels if i == 0 else self.feat_channels
            if self.dcn_on_last_conv and i == self.stacked_convs - 1:
                conv_cfg = dict(type='DCNv2')
            else:
                conv_cfg = self.conv_cfg
            self.cls_convs.append(
                ConvModule(
                    chn,
                    self.feat_channels,
                    3,
                    stride=1,
                    padding=1,
                    conv_cfg=conv_cfg,
                    norm_cfg=self.norm_cfg,
                    bias=self.conv_bias))

该方法初始化了分类卷积层。它创建了一个包含多个卷积模块的nn.ModuleList。每个卷积模块根据参数设置进行配置，在最后一个卷积层可以选择使用可变形卷积（DCNv2）。

详细解析：

1. 创建卷积模块列表：

   • self.cls_convs = nn.ModuleList()：创建一个空的nn.ModuleList，用于存储后续创建的卷积模块。这个列表将作为分类卷积层的容器。

2. 遍历堆叠卷积层数：

   • for i in range(self.stacked_convs):：这里遍历指定的堆叠卷积层数（self.stacked_convs）。这个参数通常表示在头部网络中用于分类的卷积层的数量。

3. 确定输入通道数：

   • chn = self.in_channels if i == 0 else self.feat_channels：根据当前遍历的索引i确定卷积模块的输入通道数。如果是第一个卷积层（i == 0），则输入通道数为模型的输入通道数self.in_channels；否则，输入通道数为指定的特征通道数self.feat_channels。这样可以逐步提取和转换特征，使得后续的卷积层能够处理更深层次的特征表示。

4. 设置卷积配置：

   • if self.dcn_on_last_conv and i == self.stacked_convs - 1:：检查是否在最后一个卷积层使用可变形卷积（DCNv2）。如果self.dcn_on_last_conv为真且当前是最后一个卷积层（i == self.stacked_convs - 1），则设置卷积配置为dict(type='DCNv2')，表示使用可变形卷积。
   • else:：如果不满足上述条件，则使用模型初始化时传入的卷积配置self.conv_cfg。

5. 添加卷积模块到列表：

   • self.cls_convs.append(ConvModule(...))：使用确定的输入通道数、输出通道数（固定为self.feat_channels）、卷积核大小（3x3）、步长（1）、填充（1）、卷积配置conv_cfg、归一化配置norm_cfg和偏置设置bias=self.conv_bias创建一个卷积模块，并将其添加到self.cls_convs列表中。

2.2 _init_reg_convs方法

    def _init_reg_convs(self):
        self.reg_convs = nn.ModuleList()
        for i in range(self.stacked_convs):
            chn = self.in_channels if i == 0 else self.feat_channels
            if self.dcn_on_last_conv and i == self.stacked_convs - 1:
                conv_cfg = dict(type='DCNv2')
            else:
                conv_cfg = self.conv_cfg
            self.reg_convs.append(
                ConvModule(
                    chn,
                    self.feat_channels,
                    3,
                    stride=1,
                    padding=1,
                    conv_cfg=conv_cfg,
                    norm_cfg=self.norm_cfg,
                    bias=self.conv_bias))

与_init_cls_convs类似，这个方法初始化了回归卷积层。同样创建了一个nn.ModuleList，包含多个卷积模块，配置方式与分类卷积层相似。

代码逐步解析：

1. 创建卷积模块列表：

   • self.reg_convs = nn.ModuleList()：创建一个空的nn.ModuleList对象，用于存储后续创建的边界框回归卷积模块。

2. 遍历堆叠卷积层数：

   • for i in range(self.stacked_convs):：这里遍历指定的堆叠卷积层数（由self.stacked_convs指定）。这个参数通常表示在头部网络中用于边界框回归的卷积层的数量。

3. 确定输入通道数：

   • chn = self.in_channels if i == 0 else self.feat_channels：根据当前遍历的索引i确定卷积模块的输入通道数。如果是第一个卷积层（i == 0），则输入通道数为模型的输入通道数self.in_channels；否则，输入通道数为指定的特征通道数self.feat_channels。这样可以逐步提取和转换特征，使得后续的卷积层能够处理更深层次的特征表示。

4. 设置卷积配置：

   • if self.dcn_on_last_conv and i == self.stacked_convs - 1:：检查是否在最后一个卷积层使用可变形卷积（DCNv2）。如果self.dcn_on_last_conv为真且当前是最后一个卷积层（i == self.stacked_convs - 1），则设置卷积配置为dict(type='DCNv2')，表示使用可变形卷积。
   • else:：如果不满足上述条件，则使用模型初始化时传入的卷积配置self.conv_cfg。

5. 添加卷积模块到列表：

   • self.reg_convs.append(ConvModule(...))：使用确定的输入通道数、输出通道数（固定为self.feat_channels）、卷积核大小（3x3）、步长（1）、填充（1）、卷积配置conv_cfg、归一化配置norm_cfg和偏置设置bias=self.conv_bias创建一个卷积模块，并将其添加到self.reg_convs列表中。

2.3 _init_predictor方法

    def _init_predictor(self):
        self.conv_cls = nn.Conv2d(
            self.feat_channels, self.cls_out_channels, 3, padding=1)
        self.conv_reg = nn.Conv2d(self.feat_channels, 4, 3, padding=1)

这个方法的作用是初始化预测器层。它创建了两个卷积层，分别用于分类预测和边界框回归预测。

代码逐步解析:

1. 初始化分类预测卷积层：

   • self.conv_cls = nn.Conv2d(self.feat_channels, self.cls_out_channels, 3, padding=1)：
     • nn.Conv2d是 PyTorch 中用于创建二维卷积层的类。
     • 第一个参数self.feat_channels表示输入通道数，即前面的特征提取部分输出的特征通道数。
     • 第二个参数self.cls_out_channels表示输出通道数，这个值通常与分类的类别数量有关。如果使用 sigmoid 激活函数进行二分类，输出通道数可能为类别数量；如果使用 softmax 激活函数进行多分类，输出通道数可能为类别数量加 1（包括背景类）。
     • 3表示卷积核的大小为 3x3。
     • padding=1表示在输入特征图的四周各填充一行和一列零，以保持输出特征图的尺寸与输入特征图相同（假设步长为 1）。

2. 初始化边界框回归预测卷积层：

   • self.conv_reg = nn.Conv2d(self.feat_channels, 4, 3, padding=1)：
     • 同样使用nn.Conv2d创建卷积层。
     • 输入通道数也是self.feat_channels，与分类预测卷积层相同，因为它们都是基于相同的特征进行预测。
     • 输出通道数为 4，通常这四个值分别对应边界框的左上角坐标（x, y）和宽度、高度。
     • 卷积核大小和填充与分类预测卷积层一致。

3. _load_from_state_dict方法

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
        version = local_metadata.get('version', None)
        if version is None:
            # 处理旧版本的模型状态字典键
            bbox_head_keys = [k for k in state_dict.keys() if k.startswith(prefix)]
            ori_predictor_keys = []
            new_predictor_keys = []
            for key in bbox_head_keys:
                ori_predictor_keys.append(key)
                key = key.split('.')
                if len(key) < 2:
                    conv_name = None
                elif key[1].endswith('cls'):
                    conv_name = 'conv_cls'
                elif key[1].endswith('reg'):
                    conv_name = 'conv_reg'
                elif key[1].endswith('centerness'):
                    conv_name = 'conv_centerness'
                else:
                    conv_name = None
                if conv_name is not None:
                    key[1] = conv_name
                    new_predictor_keys.append('.'.join(key))
                else:
                    ori_predictor_keys.pop(-1)
            for i in range(len(new_predictor_keys)):
                state_dict[new_predictor_keys[i]] = state_dict.pop(ori_predictor_keys[i])
        super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)

这个方法用于从模型的状态字典中加载参数。如果没有版本信息，它会处理旧版本的模型状态字典键，将旧的键名转换为新的格式，以确保兼容性。

4. forward方法

    def forward(self, x):
        return multi_apply(self.forward_single, x)[:2]

这个方法接受上游网络的特征图作为输入，通过multi_apply函数调用forward_single方法对每个特征图进行处理，并返回分类分数和边界框预测。

5. forward_single方法

    def forward_single(self, x):
        cls_feat = x
        reg_feat = x
        for cls_layer in self.cls_convs:
            cls_feat = cls_layer(cls_feat)
        cls_score = self.conv_cls(cls_feat)
        for reg_layer in self.reg_convs:
            reg_feat = reg_layer(reg_feat)
        bbox_pred = self.conv_reg(reg_feat)
        return cls_score, bbox_pred, cls_feat, reg_feat

5.1 方法功能

这个方法用于处理单个尺度的特征图，对输入的特征图进行分类和回归操作，并返回分类分数、边界框预测以及经过分类和回归卷积层处理后的特征图。

5.2 参数解释

• x：输入的特征图张量，通常是来自特征金字塔网络（FPN）中特定尺度的特征图。

5.3 代码逐步解析

1. 初始化分类和回归特征：

   • cls_feat = x：将输入特征图赋值给cls_feat，作为分类特征的初始值。
   • reg_feat = x：同样将输入特征图赋值给reg_feat，作为回归特征的初始值。

2. 分类特征处理：

   • for cls_layer in self.cls_convs:：遍历分类卷积层列表cls_convs。
   • cls_feat = cls_layer(cls_feat)：将当前的分类特征通过每个分类卷积层进行处理，逐步提取和转换分类特征。
   • cls_score = self.conv_cls(cls_feat)：使用分类预测器卷积层conv_cls对最终的分类特征进行处理，得到分类分数张量cls_score。

3. 回归特征处理：

   • for reg_layer in self.reg_convs:：遍历回归卷积层列表reg_convs。
   • reg_feat = reg_layer(reg_feat)：将当前的回归特征通过每个回归卷积层进行处理，逐步提取和转换回归特征。
   • bbox_pred = self.conv_reg(reg_feat)：使用回归预测器卷积层conv_reg对最终的回归特征进行处理，得到边界框预测张量bbox_pred。

4. 返回结果：

   • return cls_score, bbox_pred, cls_feat, reg_feat：返回一个元组，包含分类分数、边界框预测、经过分类卷积层处理后的特征图以及经过回归卷积层处理后的特征图。这个返回值可以根据具体的模型需求进行进一步处理，例如计算损失或者进行后处理操作。

5.4 方法的作用和意义

在 Anchor-Free 目标检测模型中，这个方法对于单个尺度的特征图进行独立的分类和回归处理，为模型提供了对不同尺度目标进行检测的能力。通过分别处理分类和回归特征，模型可以学习到不同的特征表示，从而更好地预测目标的类别和边界框。同时，返回的经过卷积层处理后的特征图可以为一些特定的模型（如 FCOS）提供额外的信息，用于后续的处理或分析。

6. loss_by_feat方法

    @abstractmethod
    def loss_by_feat(
            self,
            cls_scores,
            bbox_preds,
            batch_gt_instances,
            batch_img_metas,
            batch_gt_instances_ignore=None):
        raise NotImplementedError

这是一个抽象方法，需要在具体的子类中实现。它根据提取的特征计算损失，接受分类分数、边界框预测、真实标注信息等作为参数。

7. get_targets方法

    @abstractmethod
    def get_targets(self, points, batch_gt_instances):
        raise NotImplementedError

同样是一个抽象方法，用于计算回归、分类和中心度目标，接受特征图上的点和真实标注信息作为参数。

8. aug_test方法

    def aug_test(self, aug_batch_feats, aug_batch_img_metas, rescale=False):
        return self.aug_test_bboxes(aug_batch_feats, aug_batch_img_metas, rescale=rescale)

这个方法用于测试时的数据增强，通过对增强后的特征图进行处理，返回检测结果。

三、结语

mmdet/models/dense_heads/anchor_free_head.py文件实现了一个通用的 Anchor-Free 检测头，为各种具体的 Anchor-Free 模型提供了基础框架。通过初始化各种参数和层，以及定义抽象方法，使得开发者可以在这个基础上实现具体的 Anchor-Free 检测模型，提高了代码的复用性和可扩展性。在目标检测任务中，这个文件为实现高效、准确的 Anchor-Free 检测方法提供了重要的支持。