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)