From 259d429db0133779d6f7596965029e4dc3b34d8a Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 5 Sep 2024 13:46:41 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 671496532 --- ...genet_mobilenetv4_conv_medium_seg_tpu.yaml | 65 +++++++++++++++++ .../vision/modeling/backbones/mobilenet.py | 72 +++++++++++++++++++ .../modeling/backbones/mobilenet_test.py | 11 +++ 3 files changed, 148 insertions(+) create mode 100644 official/vision/configs/experiments/image_classification/imagenet_mobilenetv4_conv_medium_seg_tpu.yaml diff --git a/official/vision/configs/experiments/image_classification/imagenet_mobilenetv4_conv_medium_seg_tpu.yaml b/official/vision/configs/experiments/image_classification/imagenet_mobilenetv4_conv_medium_seg_tpu.yaml new file mode 100644 index 00000000000..d5ca5d491d2 --- /dev/null +++ b/official/vision/configs/experiments/image_classification/imagenet_mobilenetv4_conv_medium_seg_tpu.yaml @@ -0,0 +1,65 @@ +# Top-1 Acc: 77.15%/ 66.57% Val/Train @ 500 epochs +runtime: + distribution_strategy: 'tpu' + mixed_precision_dtype: 'bfloat16' +task: + model: + num_classes: 1001 + input_size: [256, 256, 3] + backbone: + mobilenet: + model_id: 'MobileNetV4ConvMediumSeg' + flat_stochastic_depth_drop_rate: false + stochastic_depth_drop_rate: 0.075 + type: 'mobilenet' + norm_activation: + norm_epsilon: 0.001 + norm_momentum: 0.997 + dropout_rate: 0.2 + losses: + l2_weight_decay: 0.0 + label_smoothing: 0.1 + train_data: + input_path: 'gs://mlcompass-data/imagenet/imagenet-2012-tfrecord/train*' + is_training: true + global_batch_size: 4096 + dtype: 'bfloat16' + aug_type: + randaug: + cutout_const: 20 + exclude_ops: ['Cutout'] + magnitude: 15 + prob_to_apply: 0.7 + type: 'randaug' + validation_data: + input_path: 'gs://mlcompass-data/imagenet/imagenet-2012-tfrecord/valid*' + is_training: false + global_batch_size: 4096 + dtype: 'bfloat16' + drop_remainder: false +trainer: + train_steps: 156000 + validation_steps: 13 + validation_interval: 312 + steps_per_loop: 312 + summary_interval: 312 + checkpoint_interval: 312 + optimizer_config: + learning_rate: + cosine: + decay_steps: 156000 + initial_learning_rate: 0.004 + name: 'CosineDecay' + type: 'cosine' + optimizer: + adamw: + exclude_from_weight_decay: ['batch_normalization'] + gradient_clip_norm: 0.0 + weight_decay_rate: 0.1 + name: 'AdamWeightDecay' + type: 'adamw' + warmup: + linear: + warmup_steps: 1560 + name: 'linear' + type: 'linear' diff --git a/official/vision/modeling/backbones/mobilenet.py b/official/vision/modeling/backbones/mobilenet.py index addf43c4d19..0929b1775b0 100644 --- a/official/vision/modeling/backbones/mobilenet.py +++ b/official/vision/modeling/backbones/mobilenet.py @@ -766,6 +766,77 @@ def uib( } +def _mnv4_conv_medium_seg_block_specs(): + """Tailored MobileNetV4ConvMedium for dense prediction, e.g. segmentation.""" + + def convbn(kernel_size, strides, filters, output=False): + return BlockSpec( + block_fn='convbn', + kernel_size=kernel_size, + filters=filters, + strides=strides, + is_output=output, + ) + + def fused_ib(kernel_size, strides, filters, output=False): + return BlockSpec( + block_fn='fused_ib', + kernel_size=kernel_size, + filters=filters, + strides=strides, + expand_ratio=4.0, + is_output=output, + ) + + def uib( + start_dw_ks, middle_dw_ks, strides, filters, expand_ratio, output=False + ): + return BlockSpec( + block_fn='uib', + start_dw_kernel_size=start_dw_ks, + middle_dw_kernel_size=middle_dw_ks, + filters=filters, + strides=strides, + expand_ratio=expand_ratio, + use_layer_scale=False, + is_output=output, + ) + + blocks = [ + convbn(3, 2, 32), + fused_ib(3, 2, 48, output=True), + # 3rd stage + uib(3, 5, 2, 80, 4.0), + uib(3, 3, 1, 80, 2.0, output=True), + # 4th stage + uib(3, 5, 2, 160, 6.0), + uib(3, 3, 1, 160, 4.0), + uib(3, 3, 1, 160, 4.0), + uib(3, 5, 1, 160, 4.0), + uib(3, 3, 1, 160, 4.0), + uib(3, 0, 1, 160, 4.0), + uib(3, 0, 1, 160, 4.0, output=True), + # 5th stage + uib(5, 5, 2, 256, 6.0), + uib(5, 5, 1, 128, 4.0), + uib(3, 5, 1, 128, 4.0), + uib(3, 5, 1, 128, 4.0), + uib(3, 0, 1, 128, 4.0), + uib(3, 5, 1, 128, 2.0), + uib(5, 5, 1, 128, 4.0), + uib(5, 0, 1, 128, 2.0, output=False), + # FC layers + convbn(1, 1, 448, output=True), + BlockSpec(block_fn='gpooling', is_output=False), + convbn(1, 1, 1280), + ] + return { + 'spec_name': 'MobileNetV4ConvMediumSeg', + 'block_spec_schema': block_spec_field_list(), + 'block_specs': block_spec_values_to_list(blocks), + } + + MNV4ConvLarge_BLOCK_SPECS = { 'spec_name': 'MobileNetV4ConvLarge', 'block_spec_schema': [ @@ -1077,6 +1148,7 @@ def mhsa_12px(): 'MobileNetV4ConvLarge': MNV4ConvLarge_BLOCK_SPECS, 'MobileNetV4HybridMedium': _mnv4_hybrid_medium_block_specs(), 'MobileNetV4HybridLarge': _mnv4_hybrid_large_block_specs(), + 'MobileNetV4ConvMediumSeg': _mnv4_conv_medium_seg_block_specs(), } diff --git a/official/vision/modeling/backbones/mobilenet_test.py b/official/vision/modeling/backbones/mobilenet_test.py index cd5a83eb6fb..a1c7a9393c7 100644 --- a/official/vision/modeling/backbones/mobilenet_test.py +++ b/official/vision/modeling/backbones/mobilenet_test.py @@ -43,6 +43,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase): 'MobileNetV4ConvLarge', 'MobileNetV4HybridMedium', 'MobileNetV4HybridLarge', + 'MobileNetV4ConvMediumSeg', ) def test_serialize_deserialize(self, model_id): # Create a network object that sets all of its config options. @@ -96,6 +97,7 @@ def test_serialize_deserialize(self, model_id): 'MobileNetV4ConvLarge', 'MobileNetV4HybridMedium', 'MobileNetV4HybridLarge', + 'MobileNetV4ConvMediumSeg', ], ) ) @@ -126,6 +128,7 @@ def test_input_specs(self, input_dim, model_id): 'MobileNetV4ConvLarge', 'MobileNetV4HybridMedium', 'MobileNetV4HybridLarge', + 'MobileNetV4ConvMediumSeg', ], [32, 224], ) @@ -153,6 +156,7 @@ def test_mobilenet_creation(self, model_id, 'MobileNetV4ConvLarge': [48, 96, 192, 512], 'MobileNetV4HybridMedium': [48, 80, 160, 256], 'MobileNetV4HybridLarge': [48, 96, 192, 512], + 'MobileNetV4ConvMediumSeg': [48, 80, 160, 448], } network = mobilenet.MobileNet(model_id=model_id, @@ -184,6 +188,7 @@ def test_mobilenet_creation(self, model_id, 'MobileNetV4ConvLarge', 'MobileNetV4HybridMedium', 'MobileNetV4HybridLarge', + 'MobileNetV4ConvMediumSeg', ], [32, 224], ) @@ -211,6 +216,7 @@ def test_mobilenet_intermediate_layers(self, model_id, input_size): 'MobileNetV4ConvLarge': [None, None, None, None], 'MobileNetV4HybridMedium': [None, None, None, None], 'MobileNetV4HybridLarge': [None, None, None, None], + 'MobileNetV4ConvMediumSeg': [None, None, None, None], } network = mobilenet.MobileNet(model_id=model_id, filter_size_scale=1.0, @@ -247,6 +253,7 @@ def test_mobilenet_intermediate_layers(self, model_id, input_size): 'MobileNetV4ConvLarge', 'MobileNetV4HybridMedium', 'MobileNetV4HybridLarge', + 'MobileNetV4ConvMediumSeg', ], [1.0, 0.75], ) @@ -285,6 +292,8 @@ def test_mobilenet_scaling(self, model_id, ('MobileNetV4HybridMedium', 0.75): 6072584, ('MobileNetV4HybridLarge', 1.0): 36648024, ('MobileNetV4HybridLarge', 0.75): 21598064, + ('MobileNetV4ConvMediumSeg', 1.0): 3787024, + ('MobileNetV4ConvMediumSeg', 0.75): 2302536, } input_size = 224 @@ -314,6 +323,7 @@ def test_mobilenet_scaling(self, model_id, 'MobileNetV4ConvLarge', 'MobileNetV4HybridMedium', 'MobileNetV4HybridLarge', + 'MobileNetV4ConvMediumSeg', ], [8, 16, 32], ) @@ -340,6 +350,7 @@ def test_mobilenet_output_stride(self, model_id, output_stride): 'MobileNetV4ConvLarge': 512, 'MobileNetV4HybridMedium': 256, 'MobileNetV4HybridLarge': 512, + 'MobileNetV4ConvMediumSeg': 448, } network = mobilenet.MobileNet(