mirror of https://github.com/hero-y/BHRL
118 lines
4.4 KiB
Python
118 lines
4.4 KiB
Python
import torch.nn as nn
|
|
from mmcv.cnn import ConvModule, bias_init_with_prob, normal_init
|
|
|
|
from ..builder import HEADS
|
|
from .anchor_head import AnchorHead
|
|
|
|
|
|
@HEADS.register_module()
|
|
class RetinaSepBNHead(AnchorHead):
|
|
""""RetinaHead with separate BN.
|
|
|
|
In RetinaHead, conv/norm layers are shared across different FPN levels,
|
|
while in RetinaSepBNHead, conv layers are shared across different FPN
|
|
levels, but BN layers are separated.
|
|
"""
|
|
|
|
def __init__(self,
|
|
num_classes,
|
|
num_ins,
|
|
in_channels,
|
|
stacked_convs=4,
|
|
conv_cfg=None,
|
|
norm_cfg=None,
|
|
init_cfg=None,
|
|
**kwargs):
|
|
assert init_cfg is None, 'To prevent abnormal initialization ' \
|
|
'behavior, init_cfg is not allowed to be set'
|
|
self.stacked_convs = stacked_convs
|
|
self.conv_cfg = conv_cfg
|
|
self.norm_cfg = norm_cfg
|
|
self.num_ins = num_ins
|
|
super(RetinaSepBNHead, self).__init__(
|
|
num_classes, in_channels, init_cfg=init_cfg, **kwargs)
|
|
|
|
def _init_layers(self):
|
|
"""Initialize layers of the head."""
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.cls_convs = nn.ModuleList()
|
|
self.reg_convs = nn.ModuleList()
|
|
for i in range(self.num_ins):
|
|
cls_convs = nn.ModuleList()
|
|
reg_convs = nn.ModuleList()
|
|
for i in range(self.stacked_convs):
|
|
chn = self.in_channels if i == 0 else self.feat_channels
|
|
cls_convs.append(
|
|
ConvModule(
|
|
chn,
|
|
self.feat_channels,
|
|
3,
|
|
stride=1,
|
|
padding=1,
|
|
conv_cfg=self.conv_cfg,
|
|
norm_cfg=self.norm_cfg))
|
|
reg_convs.append(
|
|
ConvModule(
|
|
chn,
|
|
self.feat_channels,
|
|
3,
|
|
stride=1,
|
|
padding=1,
|
|
conv_cfg=self.conv_cfg,
|
|
norm_cfg=self.norm_cfg))
|
|
self.cls_convs.append(cls_convs)
|
|
self.reg_convs.append(reg_convs)
|
|
for i in range(self.stacked_convs):
|
|
for j in range(1, self.num_ins):
|
|
self.cls_convs[j][i].conv = self.cls_convs[0][i].conv
|
|
self.reg_convs[j][i].conv = self.reg_convs[0][i].conv
|
|
self.retina_cls = nn.Conv2d(
|
|
self.feat_channels,
|
|
self.num_anchors * self.cls_out_channels,
|
|
3,
|
|
padding=1)
|
|
self.retina_reg = nn.Conv2d(
|
|
self.feat_channels, self.num_anchors * 4, 3, padding=1)
|
|
|
|
def init_weights(self):
|
|
"""Initialize weights of the head."""
|
|
super(RetinaSepBNHead, self).init_weights()
|
|
for m in self.cls_convs[0]:
|
|
normal_init(m.conv, std=0.01)
|
|
for m in self.reg_convs[0]:
|
|
normal_init(m.conv, std=0.01)
|
|
bias_cls = bias_init_with_prob(0.01)
|
|
normal_init(self.retina_cls, std=0.01, bias=bias_cls)
|
|
normal_init(self.retina_reg, std=0.01)
|
|
|
|
def forward(self, feats):
|
|
"""Forward features from the upstream network.
|
|
|
|
Args:
|
|
feats (tuple[Tensor]): Features from the upstream network, each is
|
|
a 4D-tensor.
|
|
|
|
Returns:
|
|
tuple: Usually a tuple of classification scores and bbox prediction
|
|
cls_scores (list[Tensor]): Classification scores for all scale
|
|
levels, each is a 4D-tensor, the channels number is
|
|
num_anchors * num_classes.
|
|
bbox_preds (list[Tensor]): Box energies / deltas for all scale
|
|
levels, each is a 4D-tensor, the channels number is
|
|
num_anchors * 4.
|
|
"""
|
|
cls_scores = []
|
|
bbox_preds = []
|
|
for i, x in enumerate(feats):
|
|
cls_feat = feats[i]
|
|
reg_feat = feats[i]
|
|
for cls_conv in self.cls_convs[i]:
|
|
cls_feat = cls_conv(cls_feat)
|
|
for reg_conv in self.reg_convs[i]:
|
|
reg_feat = reg_conv(reg_feat)
|
|
cls_score = self.retina_cls(cls_feat)
|
|
bbox_pred = self.retina_reg(reg_feat)
|
|
cls_scores.append(cls_score)
|
|
bbox_preds.append(bbox_pred)
|
|
return cls_scores, bbox_preds
|