No public description

PiperOrigin-RevId: 671496532

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'

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(),
 }
 
 

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(