|
|
|
|
|
import fvcore.nn.weight_init as weight_init |
|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
|
|
|
from detectron2.layers import Conv2d, FrozenBatchNorm2d, get_norm |
|
|
from detectron2.modeling import BACKBONE_REGISTRY, ResNet, ResNetBlockBase |
|
|
from detectron2.modeling.backbone.resnet import BasicStem, BottleneckBlock, DeformBottleneckBlock |
|
|
|
|
|
from .trident_conv import TridentConv |
|
|
|
|
|
__all__ = ["TridentBottleneckBlock", "make_trident_stage", "build_trident_resnet_backbone"] |
|
|
|
|
|
|
|
|
class TridentBottleneckBlock(ResNetBlockBase): |
|
|
def __init__( |
|
|
self, |
|
|
in_channels, |
|
|
out_channels, |
|
|
*, |
|
|
bottleneck_channels, |
|
|
stride=1, |
|
|
num_groups=1, |
|
|
norm="BN", |
|
|
stride_in_1x1=False, |
|
|
num_branch=3, |
|
|
dilations=(1, 2, 3), |
|
|
concat_output=False, |
|
|
test_branch_idx=-1, |
|
|
): |
|
|
""" |
|
|
Args: |
|
|
num_branch (int): the number of branches in TridentNet. |
|
|
dilations (tuple): the dilations of multiple branches in TridentNet. |
|
|
concat_output (bool): if concatenate outputs of multiple branches in TridentNet. |
|
|
Use 'True' for the last trident block. |
|
|
""" |
|
|
super().__init__(in_channels, out_channels, stride) |
|
|
|
|
|
assert num_branch == len(dilations) |
|
|
|
|
|
self.num_branch = num_branch |
|
|
self.concat_output = concat_output |
|
|
self.test_branch_idx = test_branch_idx |
|
|
|
|
|
if in_channels != out_channels: |
|
|
self.shortcut = Conv2d( |
|
|
in_channels, |
|
|
out_channels, |
|
|
kernel_size=1, |
|
|
stride=stride, |
|
|
bias=False, |
|
|
norm=get_norm(norm, out_channels), |
|
|
) |
|
|
else: |
|
|
self.shortcut = None |
|
|
|
|
|
stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride) |
|
|
|
|
|
self.conv1 = Conv2d( |
|
|
in_channels, |
|
|
bottleneck_channels, |
|
|
kernel_size=1, |
|
|
stride=stride_1x1, |
|
|
bias=False, |
|
|
norm=get_norm(norm, bottleneck_channels), |
|
|
) |
|
|
|
|
|
self.conv2 = TridentConv( |
|
|
bottleneck_channels, |
|
|
bottleneck_channels, |
|
|
kernel_size=3, |
|
|
stride=stride_3x3, |
|
|
paddings=dilations, |
|
|
bias=False, |
|
|
groups=num_groups, |
|
|
dilations=dilations, |
|
|
num_branch=num_branch, |
|
|
test_branch_idx=test_branch_idx, |
|
|
norm=get_norm(norm, bottleneck_channels), |
|
|
) |
|
|
|
|
|
self.conv3 = Conv2d( |
|
|
bottleneck_channels, |
|
|
out_channels, |
|
|
kernel_size=1, |
|
|
bias=False, |
|
|
norm=get_norm(norm, out_channels), |
|
|
) |
|
|
|
|
|
for layer in [self.conv1, self.conv2, self.conv3, self.shortcut]: |
|
|
if layer is not None: |
|
|
weight_init.c2_msra_fill(layer) |
|
|
|
|
|
def forward(self, x): |
|
|
num_branch = self.num_branch if self.training or self.test_branch_idx == -1 else 1 |
|
|
if not isinstance(x, list): |
|
|
x = [x] * num_branch |
|
|
out = [self.conv1(b) for b in x] |
|
|
out = [F.relu_(b) for b in out] |
|
|
|
|
|
out = self.conv2(out) |
|
|
out = [F.relu_(b) for b in out] |
|
|
|
|
|
out = [self.conv3(b) for b in out] |
|
|
|
|
|
if self.shortcut is not None: |
|
|
shortcut = [self.shortcut(b) for b in x] |
|
|
else: |
|
|
shortcut = x |
|
|
|
|
|
out = [out_b + shortcut_b for out_b, shortcut_b in zip(out, shortcut)] |
|
|
out = [F.relu_(b) for b in out] |
|
|
if self.concat_output: |
|
|
out = torch.cat(out) |
|
|
return out |
|
|
|
|
|
|
|
|
def make_trident_stage(block_class, num_blocks, **kwargs): |
|
|
""" |
|
|
Create a resnet stage by creating many blocks for TridentNet. |
|
|
""" |
|
|
concat_output = [False] * (num_blocks - 1) + [True] |
|
|
kwargs["concat_output_per_block"] = concat_output |
|
|
return ResNet.make_stage(block_class, num_blocks, **kwargs) |
|
|
|
|
|
|
|
|
@BACKBONE_REGISTRY.register() |
|
|
def build_trident_resnet_backbone(cfg, input_shape): |
|
|
""" |
|
|
Create a ResNet instance from config for TridentNet. |
|
|
|
|
|
Returns: |
|
|
ResNet: a :class:`ResNet` instance. |
|
|
""" |
|
|
|
|
|
norm = cfg.MODEL.RESNETS.NORM |
|
|
stem = BasicStem( |
|
|
in_channels=input_shape.channels, |
|
|
out_channels=cfg.MODEL.RESNETS.STEM_OUT_CHANNELS, |
|
|
norm=norm, |
|
|
) |
|
|
freeze_at = cfg.MODEL.BACKBONE.FREEZE_AT |
|
|
|
|
|
if freeze_at >= 1: |
|
|
for p in stem.parameters(): |
|
|
p.requires_grad = False |
|
|
stem = FrozenBatchNorm2d.convert_frozen_batchnorm(stem) |
|
|
|
|
|
|
|
|
out_features = cfg.MODEL.RESNETS.OUT_FEATURES |
|
|
depth = cfg.MODEL.RESNETS.DEPTH |
|
|
num_groups = cfg.MODEL.RESNETS.NUM_GROUPS |
|
|
width_per_group = cfg.MODEL.RESNETS.WIDTH_PER_GROUP |
|
|
bottleneck_channels = num_groups * width_per_group |
|
|
in_channels = cfg.MODEL.RESNETS.STEM_OUT_CHANNELS |
|
|
out_channels = cfg.MODEL.RESNETS.RES2_OUT_CHANNELS |
|
|
stride_in_1x1 = cfg.MODEL.RESNETS.STRIDE_IN_1X1 |
|
|
res5_dilation = cfg.MODEL.RESNETS.RES5_DILATION |
|
|
deform_on_per_stage = cfg.MODEL.RESNETS.DEFORM_ON_PER_STAGE |
|
|
deform_modulated = cfg.MODEL.RESNETS.DEFORM_MODULATED |
|
|
deform_num_groups = cfg.MODEL.RESNETS.DEFORM_NUM_GROUPS |
|
|
num_branch = cfg.MODEL.TRIDENT.NUM_BRANCH |
|
|
branch_dilations = cfg.MODEL.TRIDENT.BRANCH_DILATIONS |
|
|
trident_stage = cfg.MODEL.TRIDENT.TRIDENT_STAGE |
|
|
test_branch_idx = cfg.MODEL.TRIDENT.TEST_BRANCH_IDX |
|
|
|
|
|
assert res5_dilation in {1, 2}, "res5_dilation cannot be {}.".format(res5_dilation) |
|
|
|
|
|
num_blocks_per_stage = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3]}[depth] |
|
|
|
|
|
stages = [] |
|
|
|
|
|
res_stage_idx = {"res2": 2, "res3": 3, "res4": 4, "res5": 5} |
|
|
out_stage_idx = [res_stage_idx[f] for f in out_features] |
|
|
trident_stage_idx = res_stage_idx[trident_stage] |
|
|
max_stage_idx = max(out_stage_idx) |
|
|
for idx, stage_idx in enumerate(range(2, max_stage_idx + 1)): |
|
|
dilation = res5_dilation if stage_idx == 5 else 1 |
|
|
first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2 |
|
|
stage_kargs = { |
|
|
"num_blocks": num_blocks_per_stage[idx], |
|
|
"stride_per_block": [first_stride] + [1] * (num_blocks_per_stage[idx] - 1), |
|
|
"in_channels": in_channels, |
|
|
"bottleneck_channels": bottleneck_channels, |
|
|
"out_channels": out_channels, |
|
|
"num_groups": num_groups, |
|
|
"norm": norm, |
|
|
"stride_in_1x1": stride_in_1x1, |
|
|
"dilation": dilation, |
|
|
} |
|
|
if stage_idx == trident_stage_idx: |
|
|
assert not deform_on_per_stage[ |
|
|
idx |
|
|
], "Not support deformable conv in Trident blocks yet." |
|
|
stage_kargs["block_class"] = TridentBottleneckBlock |
|
|
stage_kargs["num_branch"] = num_branch |
|
|
stage_kargs["dilations"] = branch_dilations |
|
|
stage_kargs["test_branch_idx"] = test_branch_idx |
|
|
stage_kargs.pop("dilation") |
|
|
elif deform_on_per_stage[idx]: |
|
|
stage_kargs["block_class"] = DeformBottleneckBlock |
|
|
stage_kargs["deform_modulated"] = deform_modulated |
|
|
stage_kargs["deform_num_groups"] = deform_num_groups |
|
|
else: |
|
|
stage_kargs["block_class"] = BottleneckBlock |
|
|
blocks = ( |
|
|
make_trident_stage(**stage_kargs) |
|
|
if stage_idx == trident_stage_idx |
|
|
else ResNet.make_stage(**stage_kargs) |
|
|
) |
|
|
in_channels = out_channels |
|
|
out_channels *= 2 |
|
|
bottleneck_channels *= 2 |
|
|
|
|
|
if freeze_at >= stage_idx: |
|
|
for block in blocks: |
|
|
block.freeze() |
|
|
stages.append(blocks) |
|
|
return ResNet(stem, stages, out_features=out_features) |
|
|
|
