def custom_inception_v3(bottleneck_channel=12, bottleneck_idx=7, compressor=None, decompressor=None, short_module_names=None, **kwargs): if short_module_names is None: short_module_names = [ 'Mixed_5b', 'Mixed_5c', 'Mixed_5d', 'Mixed_6a', 'Mixed_6b', 'Mixed_6c', 'Mixed_6d', 'Mixed_6e', 'Mixed_7a', 'Mixed_7b', 'Mixed_7c', 'fc' ] if compressor is not None: compressor = get_bottleneck_processor(compressor['name'], **compressor['params']) if decompressor is not None: decompressor = get_bottleneck_processor(decompressor['name'], **decompressor['params']) bottleneck = Bottleneck4Inception3(bottleneck_channel, bottleneck_idx, compressor, decompressor) org_model = inception_v3(**kwargs) return CustomInception3(bottleneck, short_module_names, org_model)
def custom_densenet201(bottleneck_channel=12, bottleneck_idx=7, compressor=None, decompressor=None, short_feature_names=None, **kwargs): if short_feature_names is None: short_feature_names = ['denseblock3', 'transition3', 'denseblock4', 'norm5'] if compressor is not None: compressor = get_bottleneck_processor(compressor['name'], **compressor['params']) if decompressor is not None: decompressor = get_bottleneck_processor(decompressor['name'], **decompressor['params']) bottleneck = Bottleneck4DenseNets(bottleneck_channel, bottleneck_idx, compressor, decompressor) org_model = densenet201(**kwargs) return CustomDenseNet(bottleneck, short_feature_names, org_model)
def custom_resnet_fpn_backbone(backbone_name, backbone_params_config, norm_layer=misc_nn_ops.FrozenBatchNorm2d): layer1_config = backbone_params_config.get('layer1', None) layer1 = None if layer1_config is not None: compressor_config = layer1_config.get('compressor', None) compressor = None if compressor_config is None \ else get_bottleneck_processor(compressor_config['name'], **compressor_config['params']) decompressor_config = layer1_config.get('decompressor', None) decompressor = None if decompressor_config is None \ else get_bottleneck_processor(decompressor_config['name'], **decompressor_config['params']) layer1_type = layer1_config['type'] if layer1_type == 'Bottleneck4SmallResNet' and backbone_name in {'custom_resnet18', 'custom_resnet34'}: layer1 = Bottleneck4SmallResNet(layer1_config['bottleneck_channel'], compressor, decompressor) elif layer1_type == 'Bottleneck4LargeResNet'\ and backbone_name in {'custom_resnet50', 'custom_resnet101', 'custom_resnet152'}: layer1 = Bottleneck4LargeResNet(layer1_config['bottleneck_channel'], compressor, decompressor) prefix = 'custom_' start_idx = backbone_name.find(prefix) + len(prefix) org_backbone_name = backbone_name[start_idx:] if backbone_name.startswith(prefix) else backbone_name backbone = resnet.__dict__[org_backbone_name]( pretrained=backbone_params_config.get('pretrained', False), norm_layer=norm_layer ) if layer1 is not None: backbone.layer1 = layer1 trainable_layers = backbone_params_config.get('trainable_backbone_layers', 3) # select layers that wont be frozen assert 0 <= trainable_layers <= 5 layers_to_train = ['layer4', 'layer3', 'layer2', 'layer1', 'conv1'][:trainable_layers] # freeze layers only if pretrained backbone is used for name, parameter in backbone.named_parameters(): if all([not name.startswith(layer) for layer in layers_to_train]): parameter.requires_grad_(False) returned_layers = backbone_params_config.get('returned_layers', [1, 2, 3, 4]) return_layers = {f'layer{k}': str(v) for v, k in enumerate(returned_layers)} in_channels_stage2 = backbone.inplanes // 8 in_channels_list = [in_channels_stage2 * 2 ** (i - 1) for i in returned_layers] out_channels = 256 return BackboneWithFPN(backbone, return_layers, in_channels_list, out_channels)
def custom_resnet152(bottleneck_channel=12, bottleneck_idx=7, compressor=None, decompressor=None, short_module_names=None, **kwargs): if short_module_names is None: short_module_names = ['layer3', 'layer4', 'avgpool', 'fc'] if compressor is not None: compressor = get_bottleneck_processor(compressor['name'], **compressor['params']) if decompressor is not None: decompressor = get_bottleneck_processor(decompressor['name'], **decompressor['params']) bottleneck = Bottleneck4ResNet152(bottleneck_channel, bottleneck_idx, compressor, decompressor) org_model = resnet152(**kwargs) return CustomResNet(bottleneck, short_module_names, org_model)