Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import copy | |
| import torch.nn as nn | |
| from mmcv.cnn import ConvModule, Scale | |
| from mmdet.models.dense_heads.fcos_head import FCOSHead | |
| from ..builder import HEADS | |
| class NASFCOSHead(FCOSHead): | |
| """Anchor-free head used in `NASFCOS <https://arxiv.org/abs/1906.04423>`_. | |
| It is quite similar with FCOS head, except for the searched structure of | |
| classification branch and bbox regression branch, where a structure of | |
| "dconv3x3, conv3x3, dconv3x3, conv1x1" is utilized instead. | |
| """ | |
| def __init__(self, *args, init_cfg=None, **kwargs): | |
| if init_cfg is None: | |
| init_cfg = [ | |
| dict(type='Caffe2Xavier', layer=['ConvModule', 'Conv2d']), | |
| dict( | |
| type='Normal', | |
| std=0.01, | |
| override=[ | |
| dict(name='conv_reg'), | |
| dict(name='conv_centerness'), | |
| dict( | |
| name='conv_cls', | |
| type='Normal', | |
| std=0.01, | |
| bias_prob=0.01) | |
| ]), | |
| ] | |
| super(NASFCOSHead, self).__init__(*args, init_cfg=init_cfg, **kwargs) | |
| def _init_layers(self): | |
| """Initialize layers of the head.""" | |
| dconv3x3_config = dict( | |
| type='DCNv2', | |
| kernel_size=3, | |
| use_bias=True, | |
| deform_groups=2, | |
| padding=1) | |
| conv3x3_config = dict(type='Conv', kernel_size=3, padding=1) | |
| conv1x1_config = dict(type='Conv', kernel_size=1) | |
| self.arch_config = [ | |
| dconv3x3_config, conv3x3_config, dconv3x3_config, conv1x1_config | |
| ] | |
| self.cls_convs = nn.ModuleList() | |
| self.reg_convs = nn.ModuleList() | |
| for i, op_ in enumerate(self.arch_config): | |
| op = copy.deepcopy(op_) | |
| chn = self.in_channels if i == 0 else self.feat_channels | |
| assert isinstance(op, dict) | |
| use_bias = op.pop('use_bias', False) | |
| padding = op.pop('padding', 0) | |
| kernel_size = op.pop('kernel_size') | |
| module = ConvModule( | |
| chn, | |
| self.feat_channels, | |
| kernel_size, | |
| stride=1, | |
| padding=padding, | |
| norm_cfg=self.norm_cfg, | |
| bias=use_bias, | |
| conv_cfg=op) | |
| self.cls_convs.append(copy.deepcopy(module)) | |
| self.reg_convs.append(copy.deepcopy(module)) | |
| 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) | |
| self.conv_centerness = nn.Conv2d(self.feat_channels, 1, 3, padding=1) | |
| self.scales = nn.ModuleList([Scale(1.0) for _ in self.strides]) | |