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)