tf-models-nightly 2.18.0.dev20240905__py2.py3-none-any.whl → 2.18.0.dev20240906__py2.py3-none-any.whl

official/nlp/modeling/networks/mobile_bert_encoder.py

@@ -164,36 +164,6 @@ class MobileBERTEncoder(tf_keras.Model):
         attention_scores=all_attention_scores)
     super().__init__(
         inputs=self.inputs, outputs=outputs, **kwargs)
-    self._config = dict(
-        name=self.name,
-        word_vocab_size=word_vocab_size,
-        word_embed_size=word_embed_size,
-        type_vocab_size=type_vocab_size,
-        max_sequence_length=max_sequence_length,
-        num_blocks=num_blocks,
-        hidden_size=hidden_size,
-        num_attention_heads=num_attention_heads,
-        intermediate_size=intermediate_size,
-        intermediate_act_fn=intermediate_act_fn,
-        hidden_dropout_prob=hidden_dropout_prob,
-        attention_probs_dropout_prob=attention_probs_dropout_prob,
-        intra_bottleneck_size=intra_bottleneck_size,
-        initializer_range=initializer_range,
-        use_bottleneck_attention=use_bottleneck_attention,
-        key_query_shared_bottleneck=key_query_shared_bottleneck,
-        num_feedforward_networks=num_feedforward_networks,
-        normalization_type=normalization_type,
-        classifier_activation=classifier_activation,
-        input_mask_dtype=input_mask_dtype,
-        **kwargs,
-    )
-
-  def get_config(self):
-    return dict(self._config)
-
-  @classmethod
-  def from_config(cls, config):
-    return cls(**config)
 
   def get_embedding_table(self):
     return self.embedding_layer.word_embedding.embeddings

official/vision/modeling/backbones/mobilenet.py

@@ -766,6 +766,77 @@ def _mnv4_conv_medium_block_specs():
   }
 
 
+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 @@ SUPPORTED_SPECS_MAP = {
     '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 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV4ConvLarge',
               'MobileNetV4HybridMedium',
               'MobileNetV4HybridLarge',
+              'MobileNetV4ConvMediumSeg',
           ],
       )
   )
@@ -126,6 +128,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV4ConvLarge',
               'MobileNetV4HybridMedium',
               'MobileNetV4HybridLarge',
+              'MobileNetV4ConvMediumSeg',
           ],
           [32, 224],
       )
@@ -153,6 +156,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
         '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 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV4ConvLarge',
               'MobileNetV4HybridMedium',
               'MobileNetV4HybridLarge',
+              'MobileNetV4ConvMediumSeg',
           ],
           [32, 224],
       )
@@ -211,6 +216,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
         '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 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV4ConvLarge',
               'MobileNetV4HybridMedium',
               'MobileNetV4HybridLarge',
+              'MobileNetV4ConvMediumSeg',
           ],
           [1.0, 0.75],
       )
@@ -285,6 +292,8 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
         ('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 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
               'MobileNetV4ConvLarge',
               'MobileNetV4HybridMedium',
               'MobileNetV4HybridLarge',
+              'MobileNetV4ConvMediumSeg',
           ],
           [8, 16, 32],
       )
@@ -340,6 +350,7 @@ class MobileNetTest(parameterized.TestCase, tf.test.TestCase):
         'MobileNetV4ConvLarge': 512,
         'MobileNetV4HybridMedium': 256,
         'MobileNetV4HybridLarge': 512,
+        'MobileNetV4ConvMediumSeg': 448,
     }
 
     network = mobilenet.MobileNet(

official/vision/serving/semantic_segmentation.py

@@ -26,7 +26,8 @@ class SegmentationModule(export_base.ExportModule):
 
   def _build_model(self):
     input_specs = tf_keras.layers.InputSpec(
-        shape=[self._batch_size] + self._input_image_size + [3])
+        shape=[self._batch_size] + self._input_image_size + [self._num_channels]
+    )
 
     return factory.build_segmentation_model(
         input_specs=input_specs,
@@ -72,7 +73,9 @@ class SegmentationModule(export_base.ExportModule):
     if self._input_type != 'tflite':
       with tf.device('cpu:0'):
         images_spec = tf.TensorSpec(
-            shape=self._input_image_size + [3], dtype=tf.float32)
+            shape=self._input_image_size + [self._num_channels],
+            dtype=tf.float32,
+        )
         image_info_spec = tf.TensorSpec(shape=[4, 2], dtype=tf.float32)
 
         images, image_info = tf.nest.map_structure(

official/vision/serving/semantic_segmentation_test.py

@@ -49,20 +49,22 @@ class SemanticSegmentationExportTest(tf.test.TestCase, parameterized.TestCase):
         {input_type: 'serving_default'})
     tf.saved_model.save(module, save_directory, signatures=signatures)
 
-  def _get_dummy_input(self, input_type, input_image_size):
+  def _get_dummy_input(self, input_type, input_image_size, num_channels):
     """Get dummy input for the given input type."""
 
     height = input_image_size[0]
     width = input_image_size[1]
     if input_type == 'image_tensor':
-      return tf.zeros((1, height, width, 3), dtype=np.uint8)
+      return tf.zeros((1, height, width, num_channels), dtype=np.uint8)
     elif input_type == 'image_bytes':
-      image = Image.fromarray(np.zeros((height, width, 3), dtype=np.uint8))
+      image = Image.fromarray(
+          np.zeros((height, width, num_channels), dtype=np.uint8)
+      )
       byte_io = io.BytesIO()
       image.save(byte_io, 'PNG')
       return [byte_io.getvalue()]
     elif input_type == 'tf_example':
-      image_tensor = tf.zeros((height, width, 3), dtype=tf.uint8)
+      image_tensor = tf.zeros((height, width, num_channels), dtype=tf.uint8)
       encoded_jpeg = tf.image.encode_jpeg(tf.constant(image_tensor)).numpy()
       example = tf.train.Example(
           features=tf.train.Features(
@@ -73,7 +75,7 @@ class SemanticSegmentationExportTest(tf.test.TestCase, parameterized.TestCase):
               })).SerializeToString()
       return [example]
     elif input_type == 'tflite':
-      return tf.zeros((1, height, width, 3), dtype=np.float32)
+      return tf.zeros((1, height, width, num_channels), dtype=np.float32)
 
   @parameterized.parameters(
       ('image_tensor', False, [112, 112], False),
@@ -105,7 +107,7 @@ class SemanticSegmentationExportTest(tf.test.TestCase, parameterized.TestCase):
     imported = tf.saved_model.load(tmp_dir)
     segmentation_fn = imported.signatures['serving_default']
 
-    images = self._get_dummy_input(input_type, input_image_size)
+    images = self._get_dummy_input(input_type, input_image_size, num_channels=3)
     if input_type != 'tflite':
       processed_images, _ = tf.nest.map_structure(
           tf.stop_gradient,
@@ -128,6 +130,68 @@ class SemanticSegmentationExportTest(tf.test.TestCase, parameterized.TestCase):
     out = segmentation_fn(tf.constant(images))
     self.assertAllClose(out['logits'].numpy(), expected_output.numpy())
 
+  @parameterized.parameters(
+      ('image_tensor',),
+      ('tflite',),
+  )
+  def test_export_with_extra_input_channels(self, input_type):
+    tmp_dir = self.get_temp_dir()
+    num_channels = 6
+    params = exp_factory.get_exp_config('mnv2_deeplabv3_pascal')
+    params.task.init_checkpoint = None
+    params.task.model.input_size = [112, 112, num_channels]
+    params.task.export_config.rescale_output = False
+    params.task.train_data.preserve_aspect_ratio = False
+    params.task.train_data.image_feature.mean = [0.5] * num_channels
+    params.task.train_data.image_feature.stddev = [0.5] * num_channels
+    params.task.train_data.image_feature.num_channels = num_channels
+    module = semantic_segmentation.SegmentationModule(
+        params,
+        batch_size=1,
+        input_image_size=[112, 112],
+        input_type=input_type,
+        num_channels=num_channels,
+    )
+
+    self._export_from_module(module, input_type, tmp_dir)
+
+    self.assertTrue(os.path.exists(os.path.join(tmp_dir, 'saved_model.pb')))
+    self.assertTrue(
+        os.path.exists(os.path.join(tmp_dir, 'variables', 'variables.index'))
+    )
+    self.assertTrue(
+        os.path.exists(
+            os.path.join(tmp_dir, 'variables', 'variables.data-00000-of-00001')
+        )
+    )
+
+    imported = tf.saved_model.load(tmp_dir)
+    segmentation_fn = imported.signatures['serving_default']
+
+    images = self._get_dummy_input(input_type, [112, 112], num_channels)
+
+    if input_type != 'tflite':
+      processed_images, _ = tf.nest.map_structure(
+          tf.stop_gradient,
+          tf.map_fn(
+              module._build_inputs,
+              elems=tf.zeros((1, 112, 112, num_channels), dtype=tf.uint8),
+              fn_output_signature=(
+                  tf.TensorSpec(
+                      shape=[112, 112, num_channels], dtype=tf.float32
+                  ),
+                  tf.TensorSpec(shape=[4, 2], dtype=tf.float32),
+              ),
+          ),
+      )
+    else:
+      processed_images = images
+
+    logits = module.model(processed_images, training=False)['logits']
+    expected_output = tf.image.resize(logits, [112, 112], method='bilinear')
+    out = segmentation_fn(tf.constant(images))
+    self.assertAllClose(out['logits'].numpy(), expected_output.numpy())
+
   def test_export_invalid_batch_size(self):
     batch_size = 3
     tmp_dir = self.get_temp_dir()

{tf_models_nightly-2.18.0.dev20240905.dist-info → tf_models_nightly-2.18.0.dev20240906.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.18.0.dev20240905
+Version: 2.18.0.dev20240906
 Summary: TensorFlow Official Models
 Home-page: https://github.com/tensorflow/models
 Author: Google Inc.

{tf_models_nightly-2.18.0.dev20240905.dist-info → tf_models_nightly-2.18.0.dev20240906.dist-info}/RECORD

@@ -407,7 +407,7 @@ official/nlp/modeling/networks/fnet.py,sha256=DAIQCixcqDoZjGSEo1apm3pAFlm9Qlsn1I
 official/nlp/modeling/networks/fnet_test.py,sha256=mCQc67QeaPbiResoAYIUl9V3mAhp7XR4pQ5we1KcuL8,4553
 official/nlp/modeling/networks/funnel_transformer.py,sha256=5_dcVgOgJoXZ-wZT2X5w-yaNt1OTqj-whWjg73bKkSg,24127
 official/nlp/modeling/networks/funnel_transformer_test.py,sha256=XEVGmOa-sy7CTe46NQkQ18MZGdWTEB-5qYlujOndCl4,17562
-official/nlp/modeling/networks/mobile_bert_encoder.py,sha256=TNzB1ytpeGFB2zrMVZZNPzC0Y-KyelT7QuIpt0EGi7A,8662
+official/nlp/modeling/networks/mobile_bert_encoder.py,sha256=IuAw-nsQFvh4uVgroemaVM7DoclMfeL_BvfKek1qjc4,7542
 official/nlp/modeling/networks/mobile_bert_encoder_test.py,sha256=u9WqDiIuKlQuA2cEUg_InpgaXXQASGHGhxoCseyHOZI,7115
 official/nlp/modeling/networks/packed_sequence_embedding.py,sha256=Hu1hcTtfBA96dboCNcAPwe2D7x0QYWHxY4BpO5nUDUw,12810
 official/nlp/modeling/networks/packed_sequence_embedding_test.py,sha256=atjMoYpt78Uz16NXdp7l7IrJQtgcfxvkzcKDvhl5bac,5076
@@ -1063,8 +1063,8 @@ official/vision/modeling/backbones/factory.py,sha256=coJKJpPMhgM9gAc2Q7I5_CuzAaH
 official/vision/modeling/backbones/factory_test.py,sha256=7ZJRDSQ_cqJFyfqLK375V_wEqgrQpqibzNDZzNbhthU,8635
 official/vision/modeling/backbones/mobiledet.py,sha256=iEC_KbqYqUBBBwZUfRCVtqllQwK6N4T1jmiDl29B-Ys,24896
 official/vision/modeling/backbones/mobiledet_test.py,sha256=O2yfL7MSCGtKsnXr0IVUtjicrhZGGkwTXWCLtqdsL0Y,3804
-official/vision/modeling/backbones/mobilenet.py,sha256=iwUS9WSAZA6cyagw2Ld1zUlyR1MmvA9-AS7Gm4M8HZA,61286
-official/vision/modeling/backbones/mobilenet_test.py,sha256=7cl5eerD5j5UqHL8SLmpou-PjufBz8oz_cn3tqwW1vM,13057
+official/vision/modeling/backbones/mobilenet.py,sha256=WWQGODkzYlQgCAiLVsUJ2OFpM0lL6-XLa6xjyV8VFxw,63321
+official/vision/modeling/backbones/mobilenet_test.py,sha256=DALtiz7dTtDWutjdFCzEKtlLkXehJ_8AiizyetaVgIc,13565
 official/vision/modeling/backbones/resnet.py,sha256=dnYkdlYUzChGLOrQnUbwb9YJ7BDiFwgnLptks7kFb7k,16384
 official/vision/modeling/backbones/resnet_3d.py,sha256=Cq1lrlRqIg9ss_ud1iM_axW9lsTVtGYe3iA4DL9Orzk,18657
 official/vision/modeling/backbones/resnet_3d_test.py,sha256=hhCkW28UXc2peKHGgFl0MnYexFV8qTwEUkMPZ26a_MA,3799
@@ -1160,8 +1160,8 @@ official/vision/serving/export_tflite_lib.py,sha256=2AWkyEsLvMBE19m2WAa0IpyFsHV8
 official/vision/serving/export_utils.py,sha256=8mJb1MF_6kk3lbqZOZq2Lwu4A2L1KWxjnWnV_ZpYlVI,4881
 official/vision/serving/image_classification.py,sha256=wEthg6y-geVsRkAuQ1SKv-fnECMFXYuE1qR1H0yCVBA,4562
 official/vision/serving/image_classification_test.py,sha256=qZmuiQewptSQdY2iQEkx8rHjULybgADuXsQ84SjLCok,6759
-official/vision/serving/semantic_segmentation.py,sha256=1RmEOpYz-HjB-VKa3k9p3gZj3h_3ob5d-3RO9_cXfH0,4160
-official/vision/serving/semantic_segmentation_test.py,sha256=QVx12ciMA6T20d3wuKu8pe6ekEyPj8aNbKbRKxOgRik,5638
+official/vision/serving/semantic_segmentation.py,sha256=dhrJEznThcJghAO_VtAbWoR4B5x1FQZ2r5nmVHLyLdg,4221
+official/vision/serving/semantic_segmentation_test.py,sha256=XiI1YUeR1XxFQjfDFBjE-pqgHNR42Ifuz6DLQbFKmg4,7990
 official/vision/serving/video_classification.py,sha256=s37SdFoASmX0b3MubTpMdebPOvegx0Nj7yFogu5rYXE,6884
 official/vision/serving/video_classification_test.py,sha256=vx-o4y_mkgLnCrLfw-uWm6S5nrZKIcDcv_1wMEicXrE,4232
 official/vision/tasks/__init__.py,sha256=qfhL5xyDrjZez_zjw613TyciLkqtWm-INFeES7GwOPQ,995
@@ -1222,9 +1222,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=nc6A9K53OGqF25xN5St8EiWvdVbda
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
 tensorflow_models/uplift/__init__.py,sha256=mqfa55gweOdpKoaQyid4A_4u7xw__FcQeSIF0k_pYmI,999
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.18.0.dev20240905.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.18.0.dev20240905.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.18.0.dev20240905.dist-info/METADATA,sha256=l97xKfXrPnkvEh_-6VoAmIbwFiBKD9e0F-eD8oeHNzE,1432
-tf_models_nightly-2.18.0.dev20240905.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.18.0.dev20240905.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.18.0.dev20240905.dist-info/RECORD,,
+tf_models_nightly-2.18.0.dev20240906.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.18.0.dev20240906.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.18.0.dev20240906.dist-info/METADATA,sha256=GB6Gi0zUN0SQPy3AlZv5aCOqKhVlhkpZ3aL-3xN_iWk,1432
+tf_models_nightly-2.18.0.dev20240906.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.18.0.dev20240906.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.18.0.dev20240906.dist-info/RECORD,,