def build_resnet( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: # pytype: disable=annotation-type-mismatch # typed-keras """Builds ResNet backbone from a config.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 'resnet', (f'Inconsistent backbone type ' f'{backbone_type}') return ResNet( model_id=backbone_cfg.model_id, input_specs=input_specs, depth_multiplier=backbone_cfg.depth_multiplier, stem_type=backbone_cfg.stem_type, resnetd_shortcut=backbone_cfg.resnetd_shortcut, replace_stem_max_pool=backbone_cfg.replace_stem_max_pool, se_ratio=backbone_cfg.se_ratio, init_stochastic_depth_rate=backbone_cfg.stochastic_depth_drop_rate, scale_stem=backbone_cfg.scale_stem, activation=norm_activation_config.activation, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer, bn_trainable=backbone_cfg.bn_trainable)
def build_spinenet( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: """Builds SpineNet backbone from a config.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 'spinenet', (f'Inconsistent backbone type ' f'{backbone_type}') model_id = backbone_cfg.model_id if model_id not in SCALING_MAP: raise ValueError( 'SpineNet-{} is not a valid architecture.'.format(model_id)) scaling_params = SCALING_MAP[model_id] return SpineNet( input_specs=input_specs, min_level=backbone_cfg.min_level, max_level=backbone_cfg.max_level, endpoints_num_filters=scaling_params['endpoints_num_filters'], resample_alpha=scaling_params['resample_alpha'], block_repeats=scaling_params['block_repeats'], filter_size_scale=scaling_params['filter_size_scale'], init_stochastic_depth_rate=backbone_cfg.stochastic_depth_drop_rate, kernel_regularizer=l2_regularizer, activation=norm_activation_config.activation, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon)
def build_movinet( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: # pytype: disable=annotation-type-mismatch # typed-keras """Builds MoViNet backbone from a config.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() if backbone_type != 'movinet': raise ValueError(f'Inconsistent backbone type {backbone_type}') if norm_activation_config.activation is not None: raise ValueError( 'norm_activation is not used in MoViNets, but specified: %s' % norm_activation_config.activation) return Movinet( model_id=backbone_cfg.model_id, causal=backbone_cfg.causal, use_positional_encoding=backbone_cfg.use_positional_encoding, conv_type=backbone_cfg.conv_type, se_type=backbone_cfg.se_type, input_specs=input_specs, activation=backbone_cfg.activation, gating_activation=backbone_cfg.gating_activation, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer, stochastic_depth_drop_rate=backbone_cfg.stochastic_depth_drop_rate, use_external_states=backbone_cfg.use_external_states)
def build_dilated_resnet( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: # pytype: disable=annotation-type-mismatch # typed-keras """Builds ResNet backbone from a config.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 'dilated_resnet', (f'Inconsistent backbone type ' f'{backbone_type}') return DilatedResNet( model_id=backbone_cfg.model_id, output_stride=backbone_cfg.output_stride, input_specs=input_specs, stem_type=backbone_cfg.stem_type, se_ratio=backbone_cfg.se_ratio, init_stochastic_depth_rate=backbone_cfg.stochastic_depth_drop_rate, multigrid=backbone_cfg.multigrid, last_stage_repeats=backbone_cfg.last_stage_repeats, activation=norm_activation_config.activation, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer)
def build_s3d( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: # pytype: disable=annotation-type-mismatch # typed-keras """Builds S3D backbone.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 's3d' del norm_activation_config backbone = S3D( input_specs=input_specs, final_endpoint=backbone_cfg.final_endpoint, first_temporal_kernel_size=backbone_cfg.first_temporal_kernel_size, temporal_conv_start_at=backbone_cfg.temporal_conv_start_at, gating_start_at=backbone_cfg.gating_start_at, swap_pool_and_1x1x1=backbone_cfg.swap_pool_and_1x1x1, gating_style=backbone_cfg.gating_style, use_sync_bn=backbone_cfg.use_sync_bn, norm_momentum=backbone_cfg.norm_momentum, norm_epsilon=backbone_cfg.norm_epsilon, temporal_conv_type=backbone_cfg.temporal_conv_type, kernel_regularizer=l2_regularizer, depth_multiplier=backbone_cfg.depth_multiplier) return backbone
def build_darknet( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: # pytype: disable=annotation-type-mismatch # typed-keras """Builds darknet.""" backbone_config = backbone_config.get() model = Darknet( model_id=backbone_config.model_id, min_level=backbone_config.min_level, max_level=backbone_config.max_level, input_specs=input_specs, dilate=backbone_config.dilate, width_scale=backbone_config.width_scale, depth_scale=backbone_config.depth_scale, use_reorg_input=backbone_config.use_reorg_input, activation=norm_activation_config.activation, use_sync_bn=norm_activation_config.use_sync_bn, use_separable_conv=backbone_config.use_separable_conv, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer) return model
def build_movinet( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: """Builds MoViNet backbone from a config.""" l2_regularizer = l2_regularizer or tf.keras.regularizers.L2(1.5e-5) backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 'movinet', ('Inconsistent backbone type ' f'{backbone_type}') return Movinet( model_id=backbone_cfg.model_id, causal=backbone_cfg.causal, use_positional_encoding=backbone_cfg.use_positional_encoding, conv_type=backbone_cfg.conv_type, input_specs=input_specs, activation=norm_activation_config.activation, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer, stochastic_depth_drop_rate=backbone_cfg.stochastic_depth_drop_rate)
def build_resnest( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: """Builds ResNest backbone from a config.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 'resnest', (f'Inconsistent backbone type ' f'{backbone_type}') return ResNest(model_id=backbone_cfg.model_id, input_specs=input_specs, stem_type=backbone_cfg.stem_type, activation=norm_activation_config.activation)
def build_darknet( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: """Builds darknet backbone.""" backbone_cfg = backbone_config.get() model = Darknet(model_id=backbone_cfg.model_id, input_shape=input_specs, activation=norm_activation_config.activation, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer) return model
def build_revnet( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: # pytype: disable=annotation-type-mismatch # typed-keras """Builds RevNet backbone from a config.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 'revnet', (f'Inconsistent backbone type ' f'{backbone_type}') return RevNet(model_id=backbone_cfg.model_id, input_specs=input_specs, activation=norm_activation_config.activation, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer)
def build_hardnet( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: """Builds Hardnet backbone from a config.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 'hardnet', (f'Inconsistent backbone type ' f'{backbone_type}') return HardNet(model_id=backbone_cfg.model_id, input_specs=input_specs, activation=norm_activation_config.activation, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer)
def build_mobiledet( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: Optional[tf.keras.regularizers.Regularizer] = None ) -> tf.keras.Model: """Builds MobileDet backbone from a config.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 'mobiledet', (f'Inconsistent backbone type ' f'{backbone_type}') return MobileDet( model_id=backbone_cfg.model_id, filter_size_scale=backbone_cfg.filter_size_scale, input_specs=input_specs, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer)
def build_mobilenet_edgetpu(input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, **unused_kwargs) -> tf.keras.Model: """Builds MobileNetEdgeTpu backbone from a config.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 'mobilenet_edgetpu', ( f'Inconsistent backbone type ' f'{backbone_type}') if backbone_cfg.model_id in MOBILENET_EDGETPU_V2_CONFIGS: model = MobilenetEdgeTPUV2.from_name( model_name=backbone_cfg.model_id, overrides={ 'batch_norm': 'tpu', 'rescale_input': False, 'resolution': input_specs.shape[1:3], 'backbone_only': True, 'features_as_dict': True, 'dtype': 'bfloat16' }, model_weights_path=backbone_cfg.pretrained_checkpoint_path) if backbone_cfg.freeze_large_filters: freeze_large_filters(model, backbone_cfg.freeze_large_filters) return model elif backbone_cfg.model_id in MOBILENET_EDGETPU_CONFIGS: model = MobilenetEdgeTPU.from_name( model_name=backbone_cfg.model_id, overrides={ 'batch_norm': 'tpu', 'rescale_input': False, 'resolution': input_specs.shape[1:3], 'backbone_only': True, 'dtype': 'bfloat16' }, model_weights_path=backbone_cfg.pretrained_checkpoint_path) if backbone_cfg.freeze_large_filters: freeze_large_filters(model, backbone_cfg.freeze_large_filters) return model else: raise ValueError(f'Unsupported model/id type {backbone_cfg.model_id}.')
def build_efficientnet( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model: """Builds EfficientNet backbone from a config.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 'efficientnet', (f'Inconsistent backbone type ' f'{backbone_type}') return EfficientNet( model_id=backbone_cfg.model_id, input_specs=input_specs, stochastic_depth_drop_rate=backbone_cfg.stochastic_depth_drop_rate, se_ratio=backbone_cfg.se_ratio, activation=norm_activation_config.activation, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer)
def build_unet3d( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: # pytype: disable=annotation-type-mismatch # typed-keras """Builds 3D UNet backbone from a config.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 'unet_3d', (f'Inconsistent backbone type ' f'{backbone_type}') return UNet3D(model_id=backbone_cfg.model_id, input_specs=input_specs, pool_size=backbone_cfg.pool_size, base_filters=backbone_cfg.base_filters, kernel_regularizer=l2_regularizer, activation=norm_activation_config.activation, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, use_sync_bn=norm_activation_config.use_sync_bn, use_batch_normalization=backbone_cfg.use_batch_normalization)
def build_assemblenet_v1( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: Optional[tf.keras.regularizers.Regularizer] = None ) -> tf.keras.Model: """Builds assemblenet backbone.""" del l2_regularizer backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert 'assemblenet' in backbone_type assemblenet_depth = int(backbone_cfg.model_id) if assemblenet_depth not in ASSEMBLENET_SPECS: raise ValueError('Not a valid assemblenet_depth:', assemblenet_depth) model_structure, model_edge_weights = cfg.blocks_to_flat_lists( backbone_cfg.blocks) params = ASSEMBLENET_SPECS[assemblenet_depth] block_fn = functools.partial( params['block'], use_sync_bn=norm_activation_config.use_sync_bn, bn_decay=norm_activation_config.norm_momentum, bn_epsilon=norm_activation_config.norm_epsilon) backbone = AssembleNet( block_fn=block_fn, num_blocks=params['num_blocks'], num_frames=backbone_cfg.num_frames, model_structure=model_structure, input_specs=input_specs, model_edge_weights=model_edge_weights, combine_method=backbone_cfg.combine_method, use_sync_bn=norm_activation_config.use_sync_bn, bn_decay=norm_activation_config.norm_momentum, bn_epsilon=norm_activation_config.norm_epsilon) logging.info('Number of parameters in AssembleNet backbone: %f M.', backbone.count_params() / 10.**6) return backbone
def build_darknet( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: """Builds darknet.""" backbone_cfg = backbone_config.get() model = Darknet(model_id=backbone_cfg.model_id, min_level=backbone_cfg.min_level, max_level=backbone_cfg.max_level, input_specs=input_specs, dilate=backbone_cfg.dilate, width_scale=backbone_cfg.width_scale, depth_scale=backbone_cfg.depth_scale, activation=norm_activation_config.activation, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer) model.summary() return model
def build_resnet3d_rs( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: tf.keras.regularizers.Regularizer = None ) -> tf.keras.Model: """Builds ResNet-3D-RS backbone from a config.""" backbone_cfg = backbone_config.get() # Flatten configs before passing to the backbone. temporal_strides = [] temporal_kernel_sizes = [] use_self_gating = [] for i, block_spec in enumerate(backbone_cfg.block_specs): temporal_strides.append(block_spec.temporal_strides) use_self_gating.append(block_spec.use_self_gating) block_repeats_i = RESNET_SPECS[backbone_cfg.model_id][i][-1] temporal_kernel_sizes.append( list(block_spec.temporal_kernel_sizes) * block_repeats_i) return ResNet3D( model_id=backbone_cfg.model_id, temporal_strides=temporal_strides, temporal_kernel_sizes=temporal_kernel_sizes, use_self_gating=use_self_gating, input_specs=input_specs, stem_type=backbone_cfg.stem_type, stem_conv_temporal_kernel_size=backbone_cfg. stem_conv_temporal_kernel_size, stem_conv_temporal_stride=backbone_cfg.stem_conv_temporal_stride, stem_pool_temporal_stride=backbone_cfg.stem_pool_temporal_stride, init_stochastic_depth_rate=backbone_cfg.stochastic_depth_drop_rate, se_ratio=backbone_cfg.se_ratio, activation=norm_activation_config.activation, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer)
def build_hourglass( input_specs: tf.keras.layers.InputSpec, backbone_config: hyperparams.Config, norm_activation_config: hyperparams.Config, l2_regularizer: Optional[tf.keras.regularizers.Regularizer] = None ) -> tf.keras.Model: """Builds Hourglass backbone from a configuration.""" backbone_type = backbone_config.type backbone_cfg = backbone_config.get() assert backbone_type == 'hourglass', (f'Inconsistent backbone type ' f'{backbone_type}') return Hourglass( model_id=backbone_cfg.model_id, input_channel_dims=backbone_cfg.input_channel_dims, num_hourglasses=backbone_cfg.num_hourglasses, input_specs=input_specs, initial_downsample=backbone_cfg.initial_downsample, activation=norm_activation_config.activation, use_sync_bn=norm_activation_config.use_sync_bn, norm_momentum=norm_activation_config.norm_momentum, norm_epsilon=norm_activation_config.norm_epsilon, kernel_regularizer=l2_regularizer, )