def freeze(cls: Union[GenEfficientNetWrapper, MobileNetV3Wrapper],
freeze_at,
is_mobilenetv3=False):
assert -1 <= freeze_at <= len(cls.blocks)
# stem == -1, block == 0...n-1, linear == n
for i in range(-1, freeze_at + 1):
if i == -1:
# freeze stem
for module in (cls.conv_stem, cls.bn1, cls.act1):
for p in module.parameters():
p.requires_grad = False
cls.bn1 = FrozenBatchNorm2d.convert_frozen_batchnorm(cls.bn1)
elif i == len(cls.blocks):
# freeze linear
for module in (cls.conv_head, cls.act2, cls.global_pool,
cls.classifier):
for p in module.parameters():
p.requires_grad = False
if not is_mobilenetv3:
for p in cls.bn2.parameters():
p.requires_grad = False
cls.bn2 = FrozenBatchNorm2d.convert_frozen_batchnorm(cls.bn2)
else:
# freeze intermediate block
for p in cls.blocks[i].parameters():
p.requires_grad = False
cls.blocks[i] = FrozenBatchNorm2d.convert_frozen_batchnorm(
cls.blocks[i])
def build_timm_backbone(cfg, input_shape):
"""
Create a TimmNet instance from config.
Returns:
TimmNet: a :class:`TimmNet` instance.
"""
norm = cfg.MODEL.TIMMNETS.NORM
out_features = cfg.MODEL.TIMMNETS.OUT_FEATURES
model_name = cfg.MODEL.TIMMNETS.NAME
pretrained = cfg.MODEL.TIMMNETS.PRETRAINED
scriptable = cfg.MODEL.TIMMNETS.SCRIPTABLE
exportable = cfg.MODEL.TIMMNETS.EXPORTABLE
# GET MODEL BY NAME
model = timm.create_model(model_name,
pretrained,
features_only=True,
out_indices=out_features,
scriptable=scriptable,
exportable=exportable,
feature_location='expansion')
# LOAD MODEL AND CONVERT NORM
# NOTE: why I use if/else: see the strange function _load_from_state_dict in FrozenBatchNorm2d
assert norm in ["FrozenBN", "SyncBN", "BN"]
if norm == "FrozenBN":
model = FrozenBatchNorm2d.convert_frozen_batchnorm(model)
elif pretrained:
model = convert_norm_to_detectron2_format(model, norm)
else:
model = convert_norm_to_detectron2_format(model,
norm,
init_default=True)
# USE TENSORFLOW EPS, MOMENTUM defaults if model is tf pretrained
if "tf" in model_name:
model = convert_norm_eps_momentum_to_tf_defaults(model)
# FREEZE FIRST 2 LAYERS
max_block_number = int(model.feature_info[1]['module'][7:8])
# max_block_number = int(model.feature_info[1]['name'][7:8])
print(f"Freezing stem and first {max_block_number + 1} backbone blocks")
for p in model.conv_stem.parameters():
p.requires_grad = False
model.bn1 = FrozenBatchNorm2d.convert_frozen_batchnorm(model.bn1)
for block_number in range(0, max_block_number + 1):
for p in model.blocks[block_number].parameters():
p.requires_grad = False
model.blocks[
block_number] = FrozenBatchNorm2d.convert_frozen_batchnorm(
model.blocks[block_number])
return model
def _freeze_backbone(self, freeze_at):
if freeze_at < 0:
return
for stage_index in range(freeze_at):
if stage_index == 0:
m = self.stem # stage 0 is the stem
else:
m = getattr(self, "stage" + str(stage_index + 1))
for p in m.parameters():
p.requires_grad = False
FrozenBatchNorm2d.convert_frozen_batchnorm(self)
def freeze(self):
"""
Make this block not trainable.
This method sets all parameters to `requires_grad=False`,
and convert all BatchNorm layers to FrozenBatchNorm
Returns:
the block itself
"""
for p in self.parameters():
p.requires_grad = False
FrozenBatchNorm2d.convert_frozen_batchnorm(self)
return self
def freeze_at(self, stage):
stage = min(stage, len(self.out_block_inds_all_stage))
if stage < 0:
return # skip freeze, used when train from scratch
# stage == 0: freeze stem
for p in self._conv_stem.parameters():
p.requires_grad = False
FrozenBatchNorm2d.convert_frozen_batchnorm(self._bn0)
# stage >= 1: freeze blocks
if stage >= 1:
block_idx = self.out_block_inds_all_stage[stage - 1]
for i, block in enumerate(self._blocks):
if i > block_idx:
break
block.freeze()
def _freeze_backbone(self, freeze_at):
if freeze_at < 0:
return
# freeze BN layers
for m in self.modules():
if isinstance(m, nn.BatchNorm2d):
freeze_bn_params(m)
for stage_index in range(freeze_at):
if stage_index == 0:
m = self.stem # stage 0 is the stem
else:
m = getattr(self, "stage" + str(stage_index + 1))
for p in m.parameters():
p.requires_grad = False
FrozenBatchNorm2d.convert_frozen_batchnorm(self)
def _freeze_matched_bn(module, name, reg_exps, matched_names, unmatched_names):
"""
Recursive function to freeze bn layers that match specified regular expressions.
"""
res = module
# Base case: current module is a leaf node
if len(list(module.children())) == 0:
if isinstance(module, nn.modules.batchnorm._BatchNorm):
matched = False
for frozen_layers_regex in reg_exps:
if re.match(frozen_layers_regex, name):
matched = True
matched_names.append(name)
# Convert to frozen batch norm
res = FrozenBatchNorm2d.convert_frozen_batchnorm(module)
if not matched:
unmatched_names.append(name)
return res
# Recursion: current module has children
for child_name, child in module.named_children():
_name = name + "." + child_name if name != "" else child_name
new_child = _freeze_matched_bn(child, _name, reg_exps, matched_names,
unmatched_names)
if new_child is not child:
res.add_module(child_name, new_child)
return res
def __init__(self, cfg, input_shape):
super().__init__()
version = cfg.MODEL.EFFICIENTNET.VERSION
pretrained = cfg.MODEL.EFFICIENTNET.PRETRAINED
f_channels = cfg.MODEL.BIFPN.F_CHANNELS
num_fpn_layers = cfg.MODEL.BIFPN.NUM_LAYERS
backbone = EfficientNet(version,
feature_levels=(3, 4, 5),
pretrained=pretrained)
self.backbone = FrozenBatchNorm2d.convert_frozen_batchnorm(backbone)
self.backbone = backbone
out_channels = self.backbone.out_channels
self.fpn1 = FPNExtraLayers(out_channels[-1],
extra_layers=(6, 7),
f_channels=f_channels)
out_channels = [*out_channels, f_channels, f_channels]
self.fpns = nn.ModuleList([
BiFPN(out_channels, f_channels), *[
BiFPN([f_channels] * 5, f_channels)
for _ in range(num_fpn_layers - 1)
]
])
self._out_features = ["p3", "p4", "p5", "p6", "p7"]
self._out_channels = [f_channels] * 5
self._out_strides = [2**s for s in range(3, 9)]
def __init__(self, cfg, input_shape, width_mult=1.):
super(MobileNetV2, self).__init__()
self._out_features = cfg.MODEL.MOBILENETV2.OUT_FEATURES
bn = cfg.MODEL.MOBILENETV2.NORM
freeze_at = cfg.MODEL.BACKBONE.FREEZE_AT
# setting of inverted residual blocks
self.cfgs = [
# t, c, n, s
[1, 16, 1, 1, ''],
[6, 24, 2, 2, 'm2'],
[6, 32, 3, 2, 'm3'],
[6, 64, 4, 2, ''],
[6, 96, 3, 1, 'm4'],
[6, 160, 3, 2, ''],
[6, 320, 1, 1, 'm5'],
]
# building first layer
input_channel = _make_divisible(32 * width_mult,
4 if width_mult == 0.1 else 8)
layers = [conv_3x3_bn(input_shape.channels, input_channel, 2, bn)]
if freeze_at >= 1:
for p in layers[0].parameters():
p.requires_grad = False
layers[0] = FrozenBatchNorm2d.convert_frozen_batchnorm(layers[0])
# building inverted residual blocks
block = InvertedResidual
self.stage_name = ['']
self._out_feature_channels = {}
self._out_feature_strides = {}
cur_stride = 2
cur_stage = 2
for t, c, n, s, name in self.cfgs:
output_channel = _make_divisible(c * width_mult,
4 if width_mult == 0.1 else 8)
cur_stride = cur_stride * s
for i in range(n):
layers.append(
block(input_channel, output_channel, s if i == 0 else 1, t,
bn))
if cur_stage <= freeze_at:
layers[-1].freeze()
if name != '' and i == n - 1:
self._out_feature_channels[name] = output_channel
self._out_feature_strides[name] = cur_stride
cur_stage += 1
input_channel = output_channel
self.stage_name.append(name if i == n - 1 else '')
self.features = nn.Sequential(*layers)
# building last several layers
# output_channel = _make_divisible(1280 * width_mult, 4 if width_mult == 0.1 else 8) if width_mult > 1.0 else 1280
# self.conv = conv_1x1_bn(input_channel, output_channel)
# self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
# self.classifier = nn.Linear(output_channel, num_classes)
self._initialize_weights()
def build_darknet_backbone(cfg, input_shape):
"""
Create a ResNet instance from config.
Returns:
ResNet: a :class:`ResNet` instance.
"""
# need registration of new blocks/stems?
norm = cfg.MODEL.DAKRNET.NORM
stem = DarknetBasicStem(
in_channels=input_shape.channels,
out_channels=cfg.MODEL.DAKRNET.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)
# fmt: off
# OUT_FEATURES: ["s1", "s2", "s3", "s4", "s5", "gap", "linear"]
out_features = cfg.MODEL.DAKRNET.OUT_FEATURES
depth = cfg.MODEL.DAKRNET.DEPTH
in_channels = cfg.MODEL.DAKRNET.STEM_OUT_CHANNELS
num_blocks_per_stage = {53: [1, 2, 8, 8, 4]}[depth]
stages = []
# Avoid creating variables without gradients
# It consumes extra memory and may cause allreduce to fail
if "linear" in out_features:
max_stage_idx = 5
else:
out_stage_idx = [{"s1": 1, "s2": 2, "s3": 3, "s4": 4, "s5": 5}[f]
for f in out_features]
max_stage_idx = max(out_stage_idx)
for idx, stage_idx in enumerate(range(1, max_stage_idx + 1)):
stage_kargs = {
"block_class": DarknetBottleneckBlock,
"num_blocks": num_blocks_per_stage[idx],
"in_channels": in_channels,
"bottleneck_channels": in_channels,
"out_channels": in_channels * 2,
"norm": norm,
}
blocks = darknet_make_stage(**stage_kargs)
in_channels *= 2
if freeze_at >= stage_idx:
for block in blocks:
block.freeze()
stages.append(blocks)
return Darknet(stem, stages, out_features=out_features, num_classes=cfg.MODEL.DAKRNET.NUM_CLASSES)
def build_efficientnet_backbone(cfg, input_shape):
"""
Create a GenEfficientNet instance from config.
Returns:
GenEfficientNet: a :class:`GenEfficientNet` instance.
"""
freeze_at = cfg.MODEL.BACKBONE.FREEZE_AT
norm = cfg.MODEL.EFFICIENTNETS.NORM
out_features = cfg.MODEL.EFFICIENTNETS.OUT_FEATURES
model_name = cfg.MODEL.EFFICIENTNETS.NAME
pretrained = cfg.MODEL.EFFICIENTNETS.PRETRAINED
set_exportable(cfg.MODEL.EFFICIENTNETS.EXPORTABLE)
set_scriptable(cfg.MODEL.EFFICIENTNETS.SCRIPTABLE)
# GET MODEL BY NAME
is_mobilenetv3 = "mobilenetv3" in model_name
if is_mobilenetv3:
mobilenetv3._create_model = lambda model_kwargs, variant, pretrained=False: (
model_kwargs, mobilenetv3.model_urls[variant])
model_kwargs, url = getattr(mobilenetv3, model_name)()
assert not pretrained or url is not None
model = MobileNetV3Wrapper(**model_kwargs, out_features=out_features)
else:
gen_efficientnet._create_model = lambda model_kwargs, variant, pretrained=False: (
model_kwargs, gen_efficientnet.model_urls[variant])
model_kwargs, url = getattr(gen_efficientnet, model_name)()
assert not pretrained or url is not None
model = GenEfficientNetWrapper(**model_kwargs,
out_features=out_features)
# LOAD MODEL AND CONVERT NORM
# NOTE: why I use if/else: see the strange function _load_from_state_dict in FrozenBatchNorm2d
assert norm in ["FrozenBN", "SyncBN", "BN"]
if norm == "FrozenBN":
if pretrained:
load_pretrained(model, url, filter_fn=None, strict=True)
model = FrozenBatchNorm2d.convert_frozen_batchnorm(model)
else:
model = convert_norm_to_detectron2_format_and_init_default(model, norm)
if pretrained:
load_pretrained(model, url, filter_fn=None, strict=True)
# USE TENSORFLOW EPS, MOMENTUM defaults if model is tf pretrained
if "tf" in model_name:
model = convert_norm_eps_momentum_to_tf_defaults(model)
# PRUNE REDUNDANT BLOCKS
prune(model, is_mobilenetv3)
# FREEZE AT
if freeze_at >= -1:
freeze(model, freeze_at, is_mobilenetv3)
return model
def build_mobilenetv3_rw_backbone(cfg, input_shape):
""" MobileNet-V3 RW
Attn: See note in gen function for this variant.
"""
# NOTE for train set drop_rate=0.2
# pretrained model trained with non-default BN epsilon
kwargs = dict()
# kwargs['bn_eps'] = BN_EPS_TF_DEFAULT
kwargs['norm_layer'] = cfg.MODEL.MOBILENETV3.NORM
out_features = cfg.MODEL.MOBILENETV3.OUT_FEATURES
kwargs['out_features'] = out_features
kwargs['out_feature_strides'] = {
f: {
"res2": 2**2,
"res3": 2**3,
"res4": 2**4,
"res5": 2**5
}[f]
for f in out_features
}
channel_multiplier = cfg.MODEL.MOBILENETV3.CHANNEL_MULTIPLIER
freeze_at = cfg.MODEL.BACKBONE.FREEZE_AT
model = _gen_mobilenet_v3_rw('mobilenetv3_rw',
channel_multiplier=channel_multiplier,
**kwargs)
if freeze_at > 1:
for param in [
*model.conv_stem.parameters(), *model.bn1.parameters(),
*model.act1.parameters()
]:
param.requires_grad = False
model.conv_stem = FrozenBatchNorm2d.convert_frozen_batchnorm(
model.conv_stem)
model.bn1 = FrozenBatchNorm2d.convert_frozen_batchnorm(model.bn1)
model.act1 = FrozenBatchNorm2d.convert_frozen_batchnorm(model.act1)
if freeze_at >= 2:
for bidx in range(2, freeze_at + 1):
for param in model.blocks[bidx - 2].parameters():
param.requires_grad = False
model.blocks[bidx -
2] = FrozenBatchNorm2d.convert_frozen_batchnorm(
model.blocks[bidx - 2])
return model
def __init__(self, cfg, n_class=1000, input_size=224, width_mult=1.):
super(ResNet18, self).__init__()
freeze_at = cfg.MODEL.BACKBONE.FREEZE_AT
self._out_features = cfg.MODEL.RESNETS.OUT_FEATURES
self._out_feature_strides = {
"res2": 4,
"res3": 8,
"res4": 16,
"res5": 32
}
self._out_feature_channels = {
"res2": 64,
"res3": 128,
"res4": 256,
"res5": 512
}
self.stages = []
# It consumes extra memory and may cause allreduce to fail
out_stage_idx = [{
"res2": 2,
"res3": 3,
"res4": 4,
"res5": 5
}[f] for f in self._out_features]
max_stage_idx = max(out_stage_idx) # max(2,3,4,5)
self.out_stage_idx = out_stage_idx
_resnet = models.resnet18(pretrained=True)
self.conv1 = _resnet.conv1
self.bn1 = _resnet.bn1
self.relu = _resnet.relu
self.maxpool = _resnet.maxpool
if freeze_at >= 1:
for p in self.conv1.parameters():
p.requires_grad = False
for p in self.bn1.parameters():
p.requires_grad = False
self.bn1 = FrozenBatchNorm2d.convert_frozen_batchnorm(self.bn1)
self.layer1 = _resnet.layer1
self.layer2 = _resnet.layer2
self.layer3 = _resnet.layer3
self.layer4 = _resnet.layer4
self.stages = [self.layer1, self.layer2, self.layer3, self.layer4]
if freeze_at > 1:
self._freeze_backbone(freeze_at - 1)
def build_custom_backbone(cfg, input_shape, num_classes=None):
"""
Create a ResNet instance from config for TridentNet.
Returns:
ResNet: a :class:`ResNet` instance.
"""
# need registration of new blocks/stems?
norm = cfg.MODEL.CUSTOM.NORM
stem = PSRBasicStem(
in_channels=input_shape,
out_channels=cfg.MODEL.RESNETS.STEM_OUT_CHANNELS,
norm=norm,
c7x7=cfg.MODEL.CUSTOM.RESNETS.STEM.CONVF_7x7,
convf_name=cfg.MODEL.CUSTOM.RESNETS.STEM.CONVF_NAME,
rot_1x1_out=cfg.MODEL.CUSTOM.RESNETS.ROT_1x1,
noise_var=cfg.MODEL.CUSTOM.RESNETS.NOISE_VAR,
stride_psr=cfg.MODEL.CUSTOM.RESNETS.STEM.STRIDE_PSR
)
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)
# fmt: off
out_features = cfg.MODEL.RESNETS.OUT_FEATURES
depth = cfg.MODEL.RESNETS.DEPTH
width_per_group = cfg.MODEL.RESNETS.WIDTH_PER_GROUP
bottleneck_channels = width_per_group
in_channels = cfg.MODEL.RESNETS.STEM_OUT_CHANNELS
out_channels = cfg.MODEL.RESNETS.RES2_OUT_CHANNELS
conv_name = cfg.MODEL.CUSTOM.RESNETS.BLOCK.CONV_NAME
conv_1x1_rot = cfg.MODEL.CUSTOM.RESNETS.BLOCK.CONV_1x1_ROT
rot_1x1_out = cfg.MODEL.CUSTOM.RESNETS.ROT_1x1
noise_var = cfg.MODEL.CUSTOM.RESNETS.NOISE_VAR
# fmt: on
num_blocks_per_stage = {10: [1, 1, 1, 1], 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]
max_stage_idx = max(out_stage_idx)
for idx, stage_idx in enumerate(range(2, max_stage_idx + 1)):
first_stride = 1 if idx == 0 else 2
stage_kargs = {
"num_blocks": num_blocks_per_stage[idx],
"first_stride": first_stride,
"in_channels": in_channels,
"bottleneck_channels": bottleneck_channels,
"out_channels": out_channels,
"conv_name": conv_name,
"conv_1x1_rot": conv_1x1_rot,
"rot_1x1_out": rot_1x1_out,
"noise_var": noise_var,
"norm": norm
}
stage_kargs["block_class"] = PSRBottleneckBlock
blocks = (
make_custom_stage(**stage_kargs)
)
in_channels = out_channels
out_channels *= 2
bottleneck_channels *= 2
stages.append(blocks)
return ResNet(stem, stages, out_features=out_features, num_classes=num_classes)
def build_custom_resnet50_fpn_backbone(cfg, input_shape: ShapeSpec):
"""
Args:
cfg: a detectron2 CfgNode
Returns:
backbone (Backbone): backbone module, must be a subclass of :class:`Backbone`.
"""
"""
in detectron2 by default:
1. resnet base and layer1(res2) requires_grad = False
2. all the bn in resnet converted to Frozen bn
"""
# bottom_up = build_resnet_backbone(cfg, input_shape)
# in_channels_p6p7 = bottom_up.output_shape()["res5"].channels
assert "WEIGHTS" in cfg.MODEL.keys(
), "Define where to find the robust network path"
return_layers = {
"layer1": "res2",
"layer2": "res3",
"layer3": "res4",
"layer4": "res5"
}
model, _ = model_utils.make_and_restore_model(
arch="resnet50",
dataset=datasets.ImageNet(""),
resume_path=cfg.MODEL.WEIGHTS,
pytorch_pretrained=False)
cfg.MODEL.WEIGHTS = ""
resnet = model.model
freeze_level = cfg.MODEL.BACKBONE.FREEZE_AT
to_freeze_layers = [
resnet.conv1, resnet.layer1, resnet.layer2, resnet.layer3,
resnet.layer4
]
# frozen_range = [resnet.conv1, resnet.layer1]
# always freeze conv1
for module in to_freeze_layers[:freeze_level]:
for param in module.parameters():
param.requires_grad = False
resnet = FrozenBatchNorm2d.convert_frozen_batchnorm(resnet)
bottom_up = IntermediateLayerGetter(resnet, return_layers)
in_features = cfg.MODEL.FPN.IN_FEATURES
out_channels = cfg.MODEL.FPN.OUT_CHANNELS
backbone = CustomResnet50FPN(
bottom_up=bottom_up,
in_features=in_features,
out_channels=out_channels,
return_layers=return_layers,
norm=cfg.MODEL.FPN.NORM,
top_block=LastLevelMaxPool(),
fuse_type=cfg.MODEL.FPN.FUSE_TYPE,
)
backbone.__bottom_up = bottom_up
return backbone
def build_resnet_backbone_caffe_maxpool(cfg, input_shape):
"""
Create a ResNet instance from config.
Returns:
ResNet: a :class:`ResNet` instance.
"""
# need registration of new blocks/stems?
norm = cfg.MODEL.RESNETS.NORM
stem = BasicStemCaffeeMaxPool( #modified BasicStem -> BasicStemCaffeeMaxPool
in_channels=input_shape.channels,
out_channels=cfg.MODEL.RESNETS.STEM_OUT_CHANNELS,
norm=norm,
caffe_maxpool=cfg.MODEL.CAFFE_MAXPOOL, #added
)
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)
# fmt: off
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
# fmt: on
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 = []
# Avoid creating variables without gradients
# It consumes extra memory and may cause allreduce to fail
out_stage_idx = [{"res2": 2, "res3": 3, "res4": 4, "res5": 5}[f] for f in out_features]
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],
"first_stride": first_stride,
"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 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_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)
def build_gtnet_backbone_pretrain(cfg, input_channel, num_classes):
"""
Create a ResNet instance from config for TridentNet.
Returns:
ResNet: a :class:`ResNet` instance.
"""
# need registration of new blocks/stems?
norm = cfg.MODEL.RESNETS.NORM
stem = GroupBasicStem(
in_channels=input_channel,
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)
# fmt: off
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
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
# fmt: on
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],
"first_stride": first_stride,
"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:
stage_kargs["block_class"] = GTBottleneckBlock
stage_kargs["num_branch"] = num_branch
stage_kargs["dilations"] = branch_dilations
stage_kargs["test_branch_idx"] = test_branch_idx
stage_kargs.pop("dilation")
else:
stage_kargs["block_class"] = R4BottleneckBlock
blocks = (
make_grouptrident_stage(**stage_kargs)
if stage_idx == trident_stage_idx
else make_rs_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 RsNet(stem, stages, num_classes=num_classes, out_features=out_features)
def main(cfg):
setup(cfg)
dataset_names = register_datasets(cfg)
if cfg.ONLY_REGISTER_DATASETS:
return {}, cfg
LOG.info(f"Registered {len(dataset_names)} datasets:" + '\n\t' + '\n\t'.join(dataset_names))
model = build_model(cfg)
checkpoint_file = cfg.MODEL.CKPT
if checkpoint_file:
if cfg.MODEL.CKPT_REMAPPER:
if cfg.EVAL_ONLY:
LOG.warning("Running with 'EVAL_ONLY', but the checkpoint is remapped.")
checkpoint_file = CHECKPOINT_REMAPPERS[cfg.MODEL.CKPT_REMAPPER](checkpoint_file, model)
# Batchnorm2D submodules to convert to FrozenBatchnNorm2D.
modules_to_convert_frozenbb = [
(name, module) for name, module in model.named_modules() if name in cfg.MODEL.CONVERT_TO_FROZEN_BN_MODULES
]
if len(modules_to_convert_frozenbb) > 0:
module_names, modules = list(zip(*modules_to_convert_frozenbb))
LOG.info(
f"Converting BatchNorm2d -> FrozenBatchNorm2d {len(modules)} submodule(s):" + '\n\t' +
'\n\t'.join(module_names)
)
for module in modules:
FrozenBatchNorm2d.convert_frozen_batchnorm(module)
# Some checkpoints contain batchnorm layer with negative value for 'running_var'.
model = HotFixFrozenBatchNorm2d.convert_frozenbn_to_hotfix_ver(model)
Checkpointer(model).load(checkpoint_file)
if cfg.EVAL_ONLY:
assert cfg.TEST.ENABLED, "'eval-only' mode is not compatible with 'cfg.TEST.ENABLED = False'."
test_results = do_test(cfg, model, is_last=True)
if cfg.TEST.AUG.ENABLED:
test_results.update(do_test(cfg, model, is_last=True, use_tta=True))
return test_results, cfg
modules_to_freeze = cfg.MODEL.FREEZE_MODULES
if modules_to_freeze:
LOG.info(f"Freezing {len(modules_to_freeze)} submodule(s):" + '\n\t' + '\n\t'.join(modules_to_freeze))
# `requires_grad=False` must be set *before* wrapping the model with `DistributedDataParallel`
# modules_to_freeze = [x.strip() for x in cfg.MODEL.FREEZE_MODULES.split(',')]
# for module_name in cfg.MODEL.FREEZE_MODULES:
for module_name in modules_to_freeze:
freeze_submodule(model, module_name)
if comm.is_distributed():
assert d2_comm._LOCAL_PROCESS_GROUP is not None
# Convert all Batchnorm*D to nn.SyncBatchNorm.
# For faster training, the batch stats are computed over only the GPUs of the same machines (usually 8).
sync_bn_pg = d2_comm._LOCAL_PROCESS_GROUP if cfg.SOLVER.SYNCBN_USE_LOCAL_WORKERS else None
model = SyncBatchNorm.convert_sync_batchnorm(model, process_group=sync_bn_pg)
model = DistributedDataParallel(
model,
device_ids=[d2_comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=cfg.SOLVER.DDP_FIND_UNUSED_PARAMETERS
)
do_train(cfg, model)
test_results = do_test(cfg, model, is_last=True)
if cfg.TEST.AUG.ENABLED:
test_results.update(do_test(cfg, model, is_last=True, use_tta=True))
return test_results, cfg
def freeze(self):
for p in self.parameters():
p.requires_grad = False
FrozenBatchNorm2d.convert_frozen_batchnorm(self)
return self
def build_resnet_backbone(cfg, input_shape):
"""
Create a ResNet instance from config.
Returns:
ResNet: a :class:`ResNet` instance.
"""
# need registration of new blocks/stems?
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)
# fmt: off
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
# multi-stage
###############################################################################
multi_stage = cfg.MODEL.RESNETS.get("MULTI_STAGE")
if multi_stage is not None and depth != 50:
raise Exception(
"Multi-stage technique now is only available for ResNet-50")
###############################################################################
# fmt: on
assert res5_dilation in {
1, 2
}, "res5_dilation cannot be {}.".format(res5_dilation)
num_blocks_per_stage = {
18: [2, 2, 2, 2],
34: [3, 4, 6, 3],
50: [3, 4, 6, 3],
101: [3, 4, 23, 3],
152: [3, 8, 36, 3],
}[depth]
if depth in [18, 34]:
assert out_channels == 64, "Must set MODEL.RESNETS.RES2_OUT_CHANNELS = 64 for R18/R34"
assert not any(
deform_on_per_stage
), "MODEL.RESNETS.DEFORM_ON_PER_STAGE unsupported for R18/R34"
assert res5_dilation == 1, "Must set MODEL.RESNETS.RES5_DILATION = 1 for R18/R34"
assert num_groups == 1, "Must set MODEL.RESNETS.NUM_GROUPS = 1 for R18/R34"
stages = []
# Avoid creating variables without gradients
# It consumes extra memory and may cause allreduce to fail
out_stage_idx = [{
"stem": 1,
"res2": 2,
"res3": 3,
"res4": 4,
"res5": 5
}[f] for f in out_features]
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],
"first_stride": first_stride,
"in_channels": in_channels,
"out_channels": out_channels,
"norm": norm,
}
# Use BasicBlock for R18 and R34.
if depth in [18, 34]:
stage_kargs["block_class"] = BasicBlock
else:
stage_kargs["bottleneck_channels"] = bottleneck_channels
stage_kargs["stride_in_1x1"] = stride_in_1x1
stage_kargs["dilation"] = dilation
stage_kargs["num_groups"] = num_groups
if 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_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,
multi_stage=multi_stage)
def build_deepent_fused_resnet_backbone(cfg, input_shape):
"""
Create a ResNet instance from config.
Returns:
ResNet: a :class:`ResNet` instance.
"""
# need registration of new blocks/stems?
norm = cfg.MODEL.RESNETS.NORM
assert input_shape.channels == 4, f'{input_shape.channels} input channels specified, should be 4'
depth_shape = input_shape._replace(channels=1)
input_shape = input_shape._replace(channels=3)
depth_encoder = build_depth_encoder_backbone(cfg, depth_shape)
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)
# fmt: off
in_features = cfg.MODEL.RESNETS.IN_FEATURES
fuse_method = cfg.MODEL.RESNETS.FUSE_METHOD
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
# fmt: on
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 = []
# Avoid creating variables without gradients
# It consumes extra memory and may cause allreduce to fail
out_stage_idx = [{
"res2": 2,
"res3": 3,
"res4": 4,
"res5": 5
}[f] for f in out_features]
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],
"first_stride": first_stride,
"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,
"block_class": BottleneckBlock
}
blocks = 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 FusedResNet(stem,
stages,
depth_encoder=depth_encoder,
in_features=in_features,
out_features=out_features,
fuse_method=fuse_method)
def build_resnetlike_backbone(cfg, input_shape: ShapeSpec):
"""
Create a ResNet instance from config.
Returns:
ResNet: a :class:`ResNet` instance.
"""
# fmt: off
out_features = cfg.MODEL.RESNETS.OUT_FEATURES
num_classes = cfg.MODEL.BACKBONE.NUM_CLASSES if "logits" in out_features else None
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
num_blocks_per_stage = cfg.MODEL.RESNETS.BLOCKS_PER_STAGE
dilation_per_stage = cfg.MODEL.RESNETS.DILATION_PAR_STAGE
norm = cfg.MODEL.RESNETS.NORM
freeze_at = cfg.MODEL.BACKBONE.FREEZE_AT
# fmt: on
# need registration of new blocks/stems?
stem = StemBlock(
in_channels=input_shape.channels,
out_channels=cfg.MODEL.RESNETS.STEM_OUT_CHANNELS,
kernel_size=cfg.MODEL.RESNETS.STEM_KERNEL,
stride=cfg.MODEL.RESNETS.STEM_STRIDE,
pooling=nn.MaxPool2d
if cfg.MODEL.RESNETS.STEM_POOLING_ON == True else None,
norm=norm,
)
if freeze_at >= 1:
for p in stem.parameters():
p.requires_grad = False
stem = FrozenBatchNorm2d.convert_frozen_batchnorm(stem)
out_stage_idx = [{
"res2": 2,
"res3": 3,
"res4": 4,
"res5": 5,
"logits": 5
}[f] for f in out_features]
max_stage_idx = max(out_stage_idx)
blocks = []
for idx in range(2, max_stage_idx + 1):
block = {
'count': num_blocks_per_stage[idx],
'in_channels': in_channels,
'bottleneck_channels': bottleneck_channels,
'out_channels': out_channels,
'dilation': dilation_per_stage[idx]
}
in_channels = out_channels
bottleneck_channels *= 2
out_channels *= 2
blocks.append(block)
return ResNetLike(stem,
stages,
num_classes=num_classes,
out_features=out_features,
freeze_at=0)
def __init__(self, cfg, input_shape):
super(build_resnet101, self).__init__()
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)
# fmt: off
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
# fmt: on
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 = []
# Avoid creating variables without gradients
# It consumes extra memory and may cause allreduce to fail
out_stage_idx = [{
"res2": 2,
"res3": 3,
"res4": 4,
"res5": 5
}[f] for f in out_features]
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],
"first_stride": first_stride,
"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 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_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)
self.RCNN_base = ResNet(stem, stages, out_features=out_features)
self.RCNN_top = self._build_res5_block(cfg)
self.load_pretrained(cfg)
def build_efficientnetv2_backbone(cfg, input_shape):
"""
Create a GenEfficientNet instance from config.
Returns:
GenEfficientNet: a :class:`GenEfficientNet` instance.
"""
freeze_at = cfg.MODEL.BACKBONE.FREEZE_AT
norm = cfg.MODEL.EFFICIENTNETS.NORM
out_features = cfg.MODEL.EFFICIENTNETS.OUT_FEATURES
model_name = cfg.MODEL.EFFICIENTNETS.NAME
pretrained = cfg.MODEL.EFFICIENTNETS.PRETRAINED
set_exportable(cfg.MODEL.EFFICIENTNETS.EXPORTABLE)
set_scriptable(cfg.MODEL.EFFICIENTNETS.SCRIPTABLE)
# GET MODEL BY NAME
model = timm.create_model('spnasnet_100',
pretrained=pretrained,
features_only=True,
out_indices=out_features)
# LOAD MODEL AND CONVERT NORM
# NOTE: why I use if/else: see the strange function _load_from_state_dict in FrozenBatchNorm2d
assert norm in ["FrozenBN", "SyncBN", "BN"]
if norm == "FrozenBN":
model = FrozenBatchNorm2d.convert_frozen_batchnorm(model)
elif pretrained:
model = convert_norm_to_detectron2_format(model, norm)
else:
model = convert_norm_to_detectron2_format_and_init_default(model, norm)
# USE TENSORFLOW EPS, MOMENTUM defaults if model is tf pretrained
if "tf" in model_name:
model = convert_norm_eps_momentum_to_tf_defaults(model)
# PRUNE REDUNDANT BLOCKS
# FIXME PRUNE
# prune(model, is_mobilenetv3)
max_block_number = int(model.feature_info[1]['name'][7:8])
for p in model.conv_stem.parameters():
p.requires_grad = False
model.bn1 = FrozenBatchNorm2d.convert_frozen_batchnorm(model.bn1)
for block_number in range(0, max_block_number + 1):
for p in model.blocks[block_number].parameters():
p.requires_grad = False
model.blocks[
block_number] = FrozenBatchNorm2d.convert_frozen_batchnorm(
model.blocks[block_number])
# def decorate_forward(cls):
# old_forward = cls.forward
#
# def new_forward(self, x):
# x = old_forward(self, x)
#
# return x + _dummy
#
# cls.forward = new_forward
# return cls
#
# model = decorate_forward(model)
return model
def __init__(self, cfg):
super().__init__()
self.device = torch.device(cfg.MODEL.DEVICE)
self.backbone = build_backbone(cfg)
self.proposal_generator = build_proposal_generator(cfg, self.backbone.output_shape())
self.roi_heads = build_roi_heads(cfg, self.backbone.output_shape())
self.vis_period = cfg.VIS_PERIOD
self.input_format = cfg.INPUT.FORMAT
self.gt_test_proposals = cfg.MODEL.GT_TEST_PROPOSALS
self.eval_method = _EvalMethod[cfg.TEST.EVAL_METHOD.upper()]
assert len(cfg.MODEL.PIXEL_MEAN) == len(cfg.MODEL.PIXEL_STD)
num_channels = len(cfg.MODEL.PIXEL_MEAN)
pixel_mean = torch.Tensor(cfg.MODEL.PIXEL_MEAN).to(self.device).view(num_channels, 1, 1)
pixel_std = torch.Tensor(cfg.MODEL.PIXEL_STD).to(self.device).view(num_channels, 1, 1)
self.normalizer = lambda x: (x - pixel_mean) / pixel_std
from detectron2.layers import FrozenBatchNorm2d
if cfg.MODEL.BACKBONE.FREEZE:
for p in self.backbone.parameters():
p.requires_grad = False
# Also freeze batchnorm, not done in original code
# TODO(): Is this even correct, or does it just return a new module?
FrozenBatchNorm2d.convert_frozen_batchnorm(self.backbone)
logger.info('froze backbone parameters')
print('froze backbone parameters')
if cfg.MODEL.PROPOSAL_GENERATOR.FREEZE:
for p in self.proposal_generator.parameters():
p.requires_grad = False
logger.info('froze proposal generator parameters')
print('froze proposal generator parameters')
FrozenBatchNorm2d.convert_frozen_batchnorm(self.proposal_generator)
if cfg.MODEL.ROI_HEADS.FREEZE_BOX_HEAD:
for p in self.roi_heads.box_head.parameters():
p.requires_grad = False
FrozenBatchNorm2d.convert_frozen_batchnorm(self.roi_heads.box_head)
logger.info('froze roi_box_head parameters')
print('froze roi_box_head parameters')
if cfg.MODEL.ROI_MASK_HEAD.FREEZE:
for p in self.roi_heads.mask_head.parameters():
p.requires_grad = False
FrozenBatchNorm2d.convert_frozen_batchnorm(self.roi_heads.mask_head)
logger.info('froze roi mask head parameters')
if cfg.MODEL.ROI_MASK_HEAD.FREEZE_WITHOUT_PREDICTOR \
and cfg.MODEL.ROI_MASK_HEAD.FREEZE:
# Both frozen doesn't make sense and likely indicates that we forgot to
# modify a config, so better to early error.
assert False
if cfg.MODEL.ROI_MASK_HEAD.FREEZE_WITHOUT_PREDICTOR:
frozen_names = []
for n, p in self.roi_heads.mask_head.named_parameters():
if 'predictor' not in n:
p.requires_grad = False
frozen_names.append(n)
logger.info('froze roi mask head parameters without predictor')
logger.info(f'Names of frozen layers: {frozen_names}')
self.to(self.device)
def build_resnest_backbone(cfg, input_shape):
"""
Create a ResNeSt instance from config.
Returns:
ResNet: a :class:`ResNet` instance.
"""
depth = cfg.MODEL.RESNETS.DEPTH
stem_width = {50: 32, 101: 64, 152: 64, 200: 64, 269: 64}[depth]
radix = cfg.MODEL.RESNETS.RADIX
deep_stem = cfg.MODEL.RESNETS.DEEP_STEM or (radix > 1)
# need registration of new blocks/stems?
norm = cfg.MODEL.RESNETS.NORM
stem = BasicStem(
in_channels=input_shape.channels,
out_channels=cfg.MODEL.RESNETS.STEM_OUT_CHANNELS,
norm=norm,
deep_stem=deep_stem,
stem_width=stem_width,
)
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)
# fmt: off
out_features = cfg.MODEL.RESNETS.OUT_FEATURES
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
avd = cfg.MODEL.RESNETS.AVD or (radix > 1)
avg_down = cfg.MODEL.RESNETS.AVG_DOWN or (radix > 1)
bottleneck_width = cfg.MODEL.RESNETS.BOTTLENECK_WIDTH
# fmt: on
assert res5_dilation in {
1, 2
}, "res5_dilation cannot be {}.".format(res5_dilation)
num_blocks_per_stage = {
18: [2, 2, 2, 2],
34: [3, 4, 6, 3],
50: [3, 4, 6, 3],
101: [3, 4, 23, 3],
152: [3, 8, 36, 3],
200: [3, 24, 36, 3],
269: [3, 30, 48, 8],
}[depth]
if depth in [18, 34]:
assert out_channels == 64, "Must set MODEL.RESNETS.RES2_OUT_CHANNELS = 64 for R18/R34"
assert not any(
deform_on_per_stage
), "MODEL.RESNETS.DEFORM_ON_PER_STAGE unsupported for R18/R34"
assert res5_dilation == 1, "Must set MODEL.RESNETS.RES5_DILATION = 1 for R18/R34"
assert num_groups == 1, "Must set MODEL.RESNETS.NUM_GROUPS = 1 for R18/R34"
stages = []
# Avoid creating variables without gradients
# It consumes extra memory and may cause allreduce to fail
out_stage_idx = [{
"res2": 2,
"res3": 3,
"res4": 4,
"res5": 5
}[f] for f in out_features]
max_stage_idx = max(out_stage_idx)
in_channels = 2 * stem_width if deep_stem else in_channels
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],
"first_stride": first_stride,
"in_channels": in_channels,
"out_channels": out_channels,
"norm": norm,
"avd": avd,
"avg_down": avg_down,
"radix": radix,
"bottleneck_width": bottleneck_width,
}
# Use BasicBlock for R18 and R34.
if depth in [18, 34]:
stage_kargs["block_class"] = BasicBlock
else:
stage_kargs["bottleneck_channels"] = bottleneck_channels
stage_kargs["stride_in_1x1"] = stride_in_1x1
stage_kargs["dilation"] = dilation
stage_kargs["num_groups"] = num_groups
if 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_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 ResNeSt(stem, stages, out_features=out_features)