Internal change

PiperOrigin-RevId: 403523011

official/vision/beta/modeling/backbones/mobilenet.py

@@ -402,6 +402,40 @@ MNMultiAVG_BLOCK_SPECS = {
     ]
 }
 
+# Similar to MobileNetMultiAVG and used for segmentation task.
+# Reduced the filters by a factor of 2 in the last block.
+MNMultiAVG_SEG_BLOCK_SPECS = {
+    'spec_name': 'MobileNetMultiAVGSeg',
+    'block_spec_schema': [
+        'block_fn', 'kernel_size', 'strides', 'filters', 'activation',
+        'expand_ratio', 'use_normalization', 'use_bias', 'is_output'
+    ],
+    'block_specs': [
+        ('convbn', 3, 2, 32, 'relu', None, True, False, False),
+        ('invertedbottleneck', 3, 2, 32, 'relu', 3., True, False, False),
+        ('invertedbottleneck', 3, 1, 32, 'relu', 2., True, False, True),
+        ('invertedbottleneck', 5, 2, 64, 'relu', 5., True, False, False),
+        ('invertedbottleneck', 3, 1, 64, 'relu', 3., True, False, False),
+        ('invertedbottleneck', 3, 1, 64, 'relu', 2., True, False, False),
+        ('invertedbottleneck', 3, 1, 64, 'relu', 3., True, False, True),
+        ('invertedbottleneck', 5, 2, 128, 'relu', 6., True, False, False),
+        ('invertedbottleneck', 3, 1, 128, 'relu', 3., True, False, False),
+        ('invertedbottleneck', 3, 1, 128, 'relu', 3., True, False, False),
+        ('invertedbottleneck', 3, 1, 128, 'relu', 3., True, False, False),
+        ('invertedbottleneck', 3, 1, 160, 'relu', 6., True, False, False),
+        ('invertedbottleneck', 3, 1, 160, 'relu', 4., True, False, True),
+        ('invertedbottleneck', 3, 2, 192, 'relu', 6., True, False, False),
+        ('invertedbottleneck', 5, 1, 96, 'relu', 2., True, False, False),
+        ('invertedbottleneck', 5, 1, 96, 'relu', 4., True, False, False),
+        ('invertedbottleneck', 5, 1, 96, 'relu', 4., True, False, True),
+        ('convbn', 1, 1, 480, 'relu', None, True, False, False),
+        ('gpooling', None, None, None, None, None, None, None, False),
+        # Remove bias and add batch norm for the last layer to support QAT
+        # and achieve slightly better accuracy.
+        ('convbn', 1, 1, 1280, 'relu', None, True, False, False),
+    ]
+}
+
 SUPPORTED_SPECS_MAP = {
     'MobileNetV1': MNV1_BLOCK_SPECS,
     'MobileNetV2': MNV2_BLOCK_SPECS,
@@ -410,6 +444,7 @@ SUPPORTED_SPECS_MAP = {
     'MobileNetV3EdgeTPU': MNV3EdgeTPU_BLOCK_SPECS,
     'MobileNetMultiMAX': MNMultiMAX_BLOCK_SPECS,
     'MobileNetMultiAVG': MNMultiAVG_BLOCK_SPECS,
+    'MobileNetMultiAVGSeg': MNMultiAVG_SEG_BLOCK_SPECS,
 }
 
 

official/vision/beta/modeling/backbones/mobilenet_test.py

@@ -36,6 +36,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
       'MobileNetV3EdgeTPU',
       'MobileNetMultiAVG',
       'MobileNetMultiMAX',
+      'MobileNetMultiAVGSeg',
   )
   def test_serialize_deserialize(self, model_id):
     # Create a network object that sets all of its config options.
@@ -80,6 +81,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV3EdgeTPU',
               'MobileNetMultiAVG',
               'MobileNetMultiMAX',
+              'MobileNetMultiAVGSeg',
           ],
       ))
   def test_input_specs(self, input_dim, model_id):
@@ -102,6 +104,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV3EdgeTPU',
               'MobileNetMultiAVG',
               'MobileNetMultiMAX',
+              'MobileNetMultiAVGSeg',
           ],
           [32, 224],
       ))
@@ -120,6 +123,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
         'MobileNetV3EdgeTPU': [32, 48, 96, 192],
         'MobileNetMultiMAX': [32, 64, 128, 160],
         'MobileNetMultiAVG': [32, 64, 160, 192],
+        'MobileNetMultiAVGSeg': [32, 64, 160, 96],
     }
 
     network = mobilenet.MobileNet(model_id=model_id,
@@ -143,6 +147,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV3EdgeTPU',
               'MobileNetMultiAVG',
               'MobileNetMultiMAX',
+              'MobileNetMultiAVGSeg',
           ],
           [32, 224],
       ))
@@ -161,6 +166,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
         'MobileNetV3EdgeTPU': [None, None, 384, 1280],
         'MobileNetMultiMAX': [96, 128, 384, 640],
         'MobileNetMultiAVG': [64, 192, 640, 768],
+        'MobileNetMultiAVGSeg': [64, 192, 640, 384],
     }
     network = mobilenet.MobileNet(model_id=model_id,
                                   filter_size_scale=1.0,
@@ -188,6 +194,8 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV3EdgeTPU',
               'MobileNetMultiAVG',
               'MobileNetMultiMAX',
+              'MobileNetMultiMAX',
+              'MobileNetMultiAVGSeg',
           ],
           [1.0, 0.75],
       ))
@@ -209,6 +217,8 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
         ('MobileNetMultiAVG', 0.75): 2349704,
         ('MobileNetMultiMAX', 1.0): 3174560,
         ('MobileNetMultiMAX', 0.75): 2045816,
+        ('MobileNetMultiAVGSeg', 1.0): 2284000,
+        ('MobileNetMultiAVGSeg', 0.75): 1427816,
     }
 
     input_size = 224
@@ -230,6 +240,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV3EdgeTPU',
               'MobileNetMultiAVG',
               'MobileNetMultiMAX',
+              'MobileNetMultiAVGSeg',
           ],
           [8, 16, 32],
       ))
@@ -247,6 +258,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
         'MobileNetV3EdgeTPU': 192,
         'MobileNetMultiMAX': 160,
         'MobileNetMultiAVG': 192,
+        'MobileNetMultiAVGSeg': 96,
     }
 
     network = mobilenet.MobileNet(