tf-models-nightly 2.12.0.dev20230426__py2.py3-none-any.whl → 2.12.0.dev20230502__py2.py3-none-any.whl

official/core/test_utils.py

@@ -25,7 +25,7 @@ class FakeKerasModel(tf.keras.Model):
     self.dense = tf.keras.layers.Dense(4, activation=tf.nn.relu)
     self.dense2 = tf.keras.layers.Dense(4, activation=tf.nn.relu)
 
-  def call(self, inputs):
+  def call(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     return self.dense2(self.dense(inputs))
 
 

official/modeling/multitask/test_utils.py

@@ -31,7 +31,7 @@ class MockFooModel(tf.keras.Model):
     self.inputs = {"foo": tf.keras.Input(shape=(2,), dtype=tf.float32),
                    "bar": tf.keras.Input(shape=(2,), dtype=tf.float32)}
 
-  def call(self, inputs):
+  def call(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     self.add_loss(tf.zeros((1,), dtype=tf.float32))
     if "foo" in inputs:
       input_tensor = inputs["foo"]
@@ -49,7 +49,7 @@ class MockBarModel(tf.keras.Model):
     self._bar_specific_layer = tf.keras.layers.Dense(1)
     self.inputs = {"bar": tf.keras.Input(shape=(2,), dtype=tf.float32)}
 
-  def call(self, inputs):
+  def call(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     self.add_loss(tf.zeros((2,), dtype=tf.float32))
     return self._bar_specific_layer(self._share_layer(inputs["bar"]))
 

official/modeling/optimization/ema_optimizer.py

@@ -14,7 +14,7 @@
 
 """Exponential moving average optimizer."""
 
-from typing import List, Optional, Text
+from typing import List, Optional
 
 import tensorflow as tf
 
@@ -79,7 +79,7 @@ class ExponentialMovingAverage(tf.keras.optimizers.legacy.Optimizer):
                average_decay: float = 0.99,
                start_step: int = 0,
                dynamic_decay: bool = True,
-               name: Text = 'ExponentialMovingAverage',
+               name: str = 'ExponentialMovingAverage',
                **kwargs):
     """Construct a new ExponentialMovingAverage optimizer.
 
@@ -107,7 +107,7 @@ class ExponentialMovingAverage(tf.keras.optimizers.legacy.Optimizer):
     self._start_step = tf.constant(start_step, tf.float32)
     self._dynamic_decay = dynamic_decay
     self._optimizer = optimizer
-    self._track_trackable(self._optimizer, 'base_optimizer')
+    self._track_trackable(self._optimizer, 'ema_base_optimizer')
     self._average_weights = None
     self._model_weights = None
 

official/modeling/optimization/lr_schedule.py

@@ -460,10 +460,6 @@ class StepCosineDecayWithOffset(
           tf.constant(math.pi) * (global_step) /
           (init_total_steps)) + 1.0) / 2.0 + next_init_lr)
       learning_rate = cosine_learning_rate
-      tf.compat.v1.logging.info("DEBUG lr %r next lr %r", learning_rate,
-                                cosine_learning_rate)
-      tf.compat.v1.logging.info("DEBUG lr %r next lr %r inittotalstep %r",
-                                init_lr, next_init_lr, init_total_steps)
 
       for i in range(1, num_levels):
         next_init_lr = lr_levels[i]
@@ -471,9 +467,6 @@ class StepCosineDecayWithOffset(
         next_total_steps = level_total_steps[i]
         next_next_init_lr = lr_levels[i + 1] if num_levels > i + 1 else 0.
 
-        tf.compat.v1.logging.info(
-            "DEBUG step %r nilr %r nss %r nts %r nnilr %r", global_step,
-            next_init_lr, next_start_step, next_total_steps, next_next_init_lr)
         next_cosine_learning_rate = ((next_init_lr - next_next_init_lr) *
                                      (tf.cos(
                                          tf.constant(math.pi) *
@@ -482,8 +475,6 @@ class StepCosineDecayWithOffset(
                                      next_next_init_lr)
         learning_rate = tf.where(global_step >= next_start_step,
                                  next_cosine_learning_rate, learning_rate)
-        tf.compat.v1.logging.info("DEBUG lr %r next lr %r", learning_rate,
-                                  next_cosine_learning_rate)
 
     return learning_rate
 

official/nlp/modeling/layers/bigbird_attention.py

@@ -458,7 +458,7 @@ class BigBirdAttention(tf.keras.layers.MultiHeadAttention):
         to_block_size=self._to_block_size,
         rand_attn=rand_attn)
 
-  def call(self, query, value, key=None, attention_mask=None, **kwargs):
+  def call(self, query, value, key=None, attention_mask=None, **kwargs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     if not self._built_from_signature:
       self._build_from_signature(query=query, value=value, key=key)
     if key is None:

official/nlp/modeling/layers/per_dim_scale_attention.py

@@ -67,7 +67,7 @@ class PerDimScaleAttention(tf.keras.layers.MultiHeadAttention):
                                  attention_scores_dropout, value)
     return attention_output, attention_scores
 
-  def call(
+  def call(  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
       self,
       query,
       value,

official/nlp/modeling/layers/relative_attention.py

@@ -228,7 +228,7 @@ class MultiHeadRelativeAttention(tf.keras.layers.MultiHeadAttention):
                                  value)
     return attention_output
 
-  def call(self,
+  def call(self,  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
            query,
            value,
            content_attention_bias,

official/nlp/modeling/layers/spectral_normalization.py

@@ -77,7 +77,7 @@ class SpectralNormalization(tf.keras.layers.Wrapper):
     super().__init__(
         layer, name=wrapper_name, **kwargs)
 
-  def build(self, input_shape):
+  def build(self, input_shape):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     super().build(input_shape)
     self.layer.kernel._aggregation = self.aggregation  # pylint: disable=protected-access
     self._dtype = self.layer.kernel.dtype
@@ -195,7 +195,7 @@ class SpectralNormalizationConv2D(tf.keras.layers.Wrapper):
           .format(input=layer))
     super().__init__(layer, **kwargs)
 
-  def build(self, input_shape):
+  def build(self, input_shape):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     if not self.layer.built:
       self.layer.build(input_shape)
     self.layer.kernel._aggregation = self.aggregation  # pylint: disable=protected-access

official/nlp/modeling/layers/tn_expand_condense_test.py

@@ -27,7 +27,7 @@ class TNLayerTest(tf.test.TestCase, parameterized.TestCase):
   """
 
   def setUp(self):
-    super(TNLayerTest, self).setUp()
+    super().setUp()
     self.labels = np.concatenate((np.ones((50, 1)), np.zeros((50, 1))), axis=0)
 
   def _build_model(self, data, proj_multiple=2):
@@ -41,21 +41,6 @@ class TNLayerTest(tf.test.TestCase, parameterized.TestCase):
     model.add(tf.keras.layers.Dense(1, activation='sigmoid'))
     return model
 
-  @parameterized.parameters((768, 6), (1024, 2))
-  def test_keras_layer(self, input_dim, proj_multiple):
-    data = np.random.normal(size=(100, input_dim))
-    data = data.astype(np.float32)
-    tf.keras.__internal__.utils.layer_test(
-        TNExpandCondense,
-        kwargs={
-            'proj_multiplier': proj_multiple,
-            'input_shape': data.shape
-        },
-        input_shape=data.shape,
-        input_data=data,
-        expected_output_shape=(None, data.shape[-1]),
-        expected_output_dtype=data.dtype)
-
   @parameterized.parameters((768, 6), (1024, 2))
   def test_train(self, input_dim, proj_multiple):
     tf.keras.utils.set_random_seed(0)

official/nlp/modeling/models/bert_pretrainer.py

@@ -226,7 +226,7 @@ class BertPretrainerV2(tf.keras.Model):
       inputs.append(masked_lm_positions)
     self.inputs = inputs
 
-  def call(self, inputs):
+  def call(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     if isinstance(inputs, list):
       logging.warning('List inputs to BertPretrainer are discouraged.')
       inputs = dict([

official/nlp/modeling/models/electra_pretrainer.py

@@ -113,7 +113,7 @@ class ElectraPretrainer(tf.keras.Model):
         units=1,
         kernel_initializer=tf_utils.clone_initializer(mlm_initializer))
 
-  def call(self, inputs):
+  def call(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     """ELECTRA forward pass.
 
     Args:

official/nlp/modeling/models/seq2seq_transformer.py

@@ -144,7 +144,7 @@ class Seq2SeqTransformer(tf.keras.Model):
 
     return embedded_inputs, boolean_mask, input_shape, source_dtype
 
-  def call(self, inputs):
+  def call(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     """Calculate target logits or inferred target sequences.
 
     Args:

official/nlp/modeling/models/xlnet.py

@@ -117,7 +117,7 @@ class XLNetPretrainer(tf.keras.Model):
         hidden_size=self._hidden_size,
         initializer=self._initializer)
 
-  def call(self, inputs: Mapping[str, Any]):
+  def call(self, inputs: Mapping[str, Any]):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     input_word_ids = inputs['input_word_ids']
     input_type_ids = inputs['input_type_ids']
     masked_tokens = inputs['masked_tokens']
@@ -212,7 +212,7 @@ class XLNetClassifier(tf.keras.Model):
         cls_token_idx=cls_token_idx,
         name=head_name)
 
-  def call(self, inputs: Mapping[str, Any]):
+  def call(self, inputs: Mapping[str, Any]):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     input_ids = inputs['input_word_ids']
     segment_ids = inputs['input_type_ids']
     input_mask = tf.cast(inputs['input_mask'], tf.float32)
@@ -305,7 +305,7 @@ class XLNetSpanLabeler(tf.keras.Model):
         dropout_rate=self._dropout_rate,
         initializer=self._initializer)
 
-  def call(self, inputs: Mapping[str, Any]):
+  def call(self, inputs: Mapping[str, Any]):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     input_word_ids = inputs['input_word_ids']
     input_type_ids = inputs['input_type_ids']
     input_mask = inputs['input_mask']

official/projects/centernet/modeling/centernet_model.py

@@ -41,7 +41,7 @@ class CenterNetModel(tf.keras.Model):
     self._detection_generator = detection_generator
     self._head = head
 
-  def call(self,
+  def call(self,  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
            inputs: tf.Tensor,
            training: bool = None,
            **kwargs) -> Mapping[str, tf.Tensor]:

official/projects/deepmac_maskrcnn/modeling/heads/hourglass_network.py

@@ -439,7 +439,7 @@ class HourglassNetwork(tf.keras.Model):
 
     self.intermediate_relu = tf.keras.layers.ReLU()
 
-  def call(self, inputs):
+  def call(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
 
     if self.initial_downsample:
       inputs = self.downsample_input(inputs)

official/projects/roformer/roformer_attention.py

@@ -90,7 +90,7 @@ class RoformerAttention(tf.keras.layers.MultiHeadAttention):
                                ...] + k2 * self.k_sin_vec[:, 0:k_len, ...]
     return ret_q, ret_w, v
 
-  def call(self,
+  def call(self,  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
            query,
            value,
            key=None,

official/projects/teams/teams_pretrainer.py

@@ -299,7 +299,7 @@ class TeamsPretrainer(tf.keras.Model):
         output=output_type,
         name='discriminator_mws')
 
-  def call(self, inputs):
+  def call(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     """TEAMS forward pass.
 
     Args:

official/projects/yolo/serving/export_module_factory.py

@@ -21,6 +21,7 @@ import tensorflow as tf
 from official.core import config_definitions as cfg
 from official.core import export_base
 from official.projects.yolo.configs.yolo import YoloTask
+from official.projects.yolo.configs.yolov7 import YoloV7Task
 from official.projects.yolo.modeling import factory as yolo_factory
 from official.projects.yolo.modeling.backbones import darknet  # pylint: disable=unused-import
 from official.projects.yolo.modeling.decoders import yolo_decoder  # pylint: disable=unused-import
@@ -163,10 +164,16 @@ def create_yolo_export_module(
       input_type, batch_size, input_image_size, num_channels, input_name)
   input_specs = tf.keras.layers.InputSpec(shape=[batch_size] +
                                           input_image_size + [num_channels])
-  model, _ = yolo_factory.build_yolo(
-      input_specs=input_specs,
-      model_config=params.task.model,
-      l2_regularization=None)
+  if isinstance(params.task, YoloTask):
+    model, _ = yolo_factory.build_yolo(
+        input_specs=input_specs,
+        model_config=params.task.model,
+        l2_regularization=None)
+  elif isinstance(params.task, YoloV7Task):
+    model = yolo_factory.build_yolov7(
+        input_specs=input_specs,
+        model_config=params.task.model,
+        l2_regularization=None)
 
   def preprocess_fn(inputs):
     image_tensor = export_utils.parse_image(inputs, input_type,
@@ -247,7 +254,7 @@ def get_export_module(params: cfg.ExperimentConfig,
                                                         input_image_size,
                                                         num_channels,
                                                         input_name)
-  elif isinstance(params.task, YoloTask):
+  elif isinstance(params.task, (YoloTask, YoloV7Task)):
     export_module = create_yolo_export_module(params, input_type, batch_size,
                                               input_image_size, num_channels,
                                               input_name)

official/projects/yolo/serving/export_saved_model.py

@@ -38,9 +38,8 @@ from absl import flags
 
 from official.core import exp_factory
 from official.modeling import hyperparams
-from official.projects.yolo.configs import yolo as cfg  # pylint: disable=unused-import
+from official.projects.yolo.common import registry_imports  # pylint: disable=unused-import
 from official.projects.yolo.serving import export_module_factory
-from official.projects.yolo.tasks import yolo as task  # pylint: disable=unused-import
 from official.vision.serving import export_saved_model_lib
 
 FLAGS = flags.FLAGS

official/utils/testing/mock_task.py

@@ -31,7 +31,7 @@ class MockModel(tf.keras.Model):
     super().__init__()
     self.network = network
 
-  def call(self, inputs):
+  def call(self, inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     outputs = self.network(inputs)
     self.add_loss(tf.reduce_mean(outputs))
     return outputs

official/vision/configs/backbones.py

@@ -50,6 +50,8 @@ class VisionTransformer(hyperparams.Config):
 
   # Adding Layerscale to each Encoder block https://arxiv.org/abs/2204.07118
   layer_scale_init_value: float = 0.0
+  # Transformer encoder spatial partition dimensions.
+  transformer_partition_dims: Optional[Tuple[int, int, int, int]] = None
 
 
 @dataclasses.dataclass

official/vision/modeling/backbones/vit.py

@@ -122,6 +122,7 @@ class Encoder(layers.Layer):
                pos_embed_origin_shape=None,
                pos_embed_target_shape=None,
                layer_scale_init_value=0.0,
+               transformer_partition_dims=None,
                **kwargs):
     super().__init__(**kwargs)
     self._num_layers = num_layers
@@ -137,6 +138,7 @@ class Encoder(layers.Layer):
     self._pos_embed_origin_shape = pos_embed_origin_shape
     self._pos_embed_target_shape = pos_embed_target_shape
     self._layer_scale_init_value = layer_scale_init_value
+    self._transformer_partition_dims = transformer_partition_dims
 
   def build(self, input_shape):
     if self._add_pos_embed:
@@ -163,7 +165,8 @@ class Encoder(layers.Layer):
           stochastic_depth_drop_rate=nn_layers.get_stochastic_depth_rate(
               self._init_stochastic_depth_rate, i + 1, self._num_layers),
           norm_epsilon=1e-6,
-          layer_scale_init_value=self._layer_scale_init_value,)
+          layer_scale_init_value=self._layer_scale_init_value,
+          transformer_partition_dims=self._transformer_partition_dims)
       self._encoder_layers.append(encoder_layer)
     self._norm = layers.LayerNormalization(epsilon=1e-6)
     super().build(input_shape)
@@ -195,6 +198,7 @@ class Encoder(layers.Layer):
         'pos_embed_origin_shape': self._pos_embed_origin_shape,
         'pos_embed_target_shape': self._pos_embed_target_shape,
         'layer_scale_init_value': self._layer_scale_init_value,
+        'transformer_partition_dims': self._transformer_partition_dims,
     }
     config.update(updates)
     return config
@@ -203,24 +207,27 @@ class Encoder(layers.Layer):
 class VisionTransformer(tf.keras.Model):
   """Class to build VisionTransformer family model."""
 
-  def __init__(self,
-               mlp_dim=3072,
-               num_heads=12,
-               num_layers=12,
-               attention_dropout_rate=0.0,
-               dropout_rate=0.1,
-               init_stochastic_depth_rate=0.0,
-               input_specs=layers.InputSpec(shape=[None, None, None, 3]),
-               patch_size=16,
-               hidden_size=768,
-               representation_size=0,
-               pooler='token',
-               kernel_regularizer=None,
-               original_init: bool = True,
-               output_encoded_tokens: bool = True,
-               output_2d_feature_maps: bool = False,
-               pos_embed_shape: Optional[Tuple[int, int]] = None,
-               layer_scale_init_value: float = 0.0):
+  def __init__(
+      self,
+      mlp_dim=3072,
+      num_heads=12,
+      num_layers=12,
+      attention_dropout_rate=0.0,
+      dropout_rate=0.1,
+      init_stochastic_depth_rate=0.0,
+      input_specs=layers.InputSpec(shape=[None, None, None, 3]),
+      patch_size=16,
+      hidden_size=768,
+      representation_size=0,
+      pooler='token',
+      kernel_regularizer=None,
+      original_init: bool = True,
+      output_encoded_tokens: bool = True,
+      output_2d_feature_maps: bool = False,
+      pos_embed_shape: Optional[Tuple[int, int]] = None,
+      layer_scale_init_value: float = 0.0,
+      transformer_partition_dims: Optional[Tuple[int, int, int, int]] = None,
+  ):
     """VisionTransformer initialization function."""
     self._mlp_dim = mlp_dim
     self._num_heads = num_heads
@@ -368,4 +375,5 @@ def build_vit(input_specs,
       output_encoded_tokens=backbone_cfg.output_encoded_tokens,
       output_2d_feature_maps=backbone_cfg.output_2d_feature_maps,
       layer_scale_init_value=backbone_cfg.layer_scale_init_value,
-      pos_embed_shape=backbone_cfg.pos_embed_shape)
+      pos_embed_shape=backbone_cfg.pos_embed_shape,
+      transformer_partition_dims=backbone_cfg.transformer_partition_dims)

official/vision/modeling/heads/segmentation_heads.py

@@ -167,7 +167,7 @@ class MaskScoring(tf.keras.Model):
 
     super(MaskScoring, self).build(input_shape)
 
-  def call(self, inputs: tf.Tensor, training: bool = None):
+  def call(self, inputs: tf.Tensor, training: bool = None):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
     """Forward pass mask scoring head.
 
     Args:

official/vision/modeling/layers/nn_blocks.py

@@ -1578,6 +1578,7 @@ class TransformerEncoderBlock(nlp_modeling.layers.TransformerEncoderBlock):
       *args,
       stochastic_depth_drop_rate=0.0,
       layer_scale_init_value=0.0,
+      transformer_partition_dims=None,
       max_attention_inference_parallelism=None,
       **kwargs
   ):
@@ -1587,6 +1588,7 @@ class TransformerEncoderBlock(nlp_modeling.layers.TransformerEncoderBlock):
       *args: positional arguments passed to super().__init__.
       stochastic_depth_drop_rate: the drop rate for the stochastic depth layer.
       layer_scale_init_value:
+      transformer_partition_dims: transformer spatial partition dimenstions.
       max_attention_inference_parallelism: the number of examples to run in
         parallel in the attention blocks during inference. Set this limit to
         reduce the peak memory usage. If None, use vectorized operations to run
@@ -1596,11 +1598,14 @@ class TransformerEncoderBlock(nlp_modeling.layers.TransformerEncoderBlock):
     super().__init__(*args, **kwargs)
     self._stochastic_depth_drop_rate = stochastic_depth_drop_rate
     self._layer_scale_init_value = layer_scale_init_value
+    self._transformer_partition_dims = transformer_partition_dims
     self._max_attention_inference_parallelism = (
         max_attention_inference_parallelism
     )
 
   def build(self, input_shape):
+    super().build(input_shape)
+
     if self._stochastic_depth_drop_rate:
       self._stochastic_depth = nn_layers.StochasticDepth(
           self._stochastic_depth_drop_rate)
@@ -1615,22 +1620,32 @@ class TransformerEncoderBlock(nlp_modeling.layers.TransformerEncoderBlock):
     else:
       self._layer_scale_attn = lambda x, *args, **kwargs: tf.identity(x)
       self._layer_scale_mlp = lambda x, *args, **kwargs: tf.identity(x)
-    super().build(input_shape)
 
-    if self._max_attention_inference_parallelism is not None:
-      attention_layer_config = self._attention_layer.get_config()
-      self._attention_layer = nn_layers.MultiHeadAttention.from_config({
-          **attention_layer_config,
-          'max_inference_parallelism': (
-              self._max_attention_inference_parallelism
-          ),
-      })
+    self._attention_layer = nn_layers.MultiHeadAttention(
+        num_heads=self._num_heads,
+        key_dim=self._key_dim,
+        value_dim=self._value_dim,
+        dropout=self._attention_dropout_rate,
+        use_bias=self._use_bias,
+        kernel_initializer=self._attention_initializer,
+        bias_initializer=tf_utils.clone_initializer(self._bias_initializer),
+        attention_axes=self._attention_axes,
+        output_shape=self._output_last_dim,
+        bias_regularizer=self._bias_regularizer,
+        activity_regularizer=self._activity_regularizer,
+        kernel_constraint=self._kernel_constraint,
+        bias_constraint=self._bias_constraint,
+        max_inference_parallelism=self._max_attention_inference_parallelism,
+        partition_dims=self._transformer_partition_dims,
+        name='self_attention',
+    )
 
   def get_config(self):
     config = super().get_config()
     config.update({
         'stochastic_depth_drop_rate': self._stochastic_depth_drop_rate,
         'layer_scale_init_value': self._layer_scale_init_value,
+        'transformer_partition_dims': self._transformer_partition_dims,
         'max_attention_inference_parallelism': (
             self._max_attention_inference_parallelism
         ),

official/vision/modeling/layers/nn_layers.py

@@ -1291,23 +1291,30 @@ class MultiHeadAttention(tf.keras.layers.MultiHeadAttention):
   """
 
   def __init__(
-      self, *args, max_inference_parallelism: Optional[int] = None, **kwargs
+      self,
+      *args,
+      partition_dims: Optional[Tuple[int, int, int, int]] = None,
+      max_inference_parallelism: Optional[int] = None,
+      **kwargs,
   ):
     """Initializes MultiHeadAttention.
 
     Args:
       *args: Positional arguments passed to super().__init__.
+      partition_dims: Spatial partition dimensions.
       max_inference_parallelism: The number of examples to run in parallel
         during inference. Set this limit to reduce the peak memory usage. If
         None, use vectorized operations to run the whole batch in parallel.
       **kwargs: Keyword arguments passed to super().__init__.
     """
     super().__init__(*args, **kwargs)
+    self._partition_dims = partition_dims
     self._max_inference_parallelism = max_inference_parallelism
 
   def get_config(self):
     config = super().get_config()
     config.update({
+        'partition_dims': self._partition_dims,
         'max_inference_parallelism': self._max_inference_parallelism,
     })
     return config
@@ -1336,6 +1343,16 @@ class MultiHeadAttention(tf.keras.layers.MultiHeadAttention):
       attention_output: Multi-headed outputs of attention computation.
       attention_scores: Multi-headed attention weights.
     """
+    if self._partition_dims is not None:
+      strategy = tf.distribute.get_strategy()
+      # `query` = [B, T, N ,H]
+      query = strategy.experimental_split_to_logical_devices(
+          query, self._partition_dims)
+      key = strategy.experimental_split_to_logical_devices(
+          key, self._partition_dims)
+      value = strategy.experimental_split_to_logical_devices(
+          value, self._partition_dims)
+
     batch_size = query.get_shape().as_list()[0]  # None if dynamic.
 
     if (

official/vision/modeling/maskrcnn_model.py

@@ -141,7 +141,7 @@ class MaskRCNNModel(tf.keras.Model):
         [30.0, 30.0, 15.0, 15.0],
     ]
 
-  def call(
+  def call(  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
       self,
       images: tf.Tensor,
       image_shape: tf.Tensor,

official/vision/modeling/segmentation_model.py

@@ -59,7 +59,7 @@ class SegmentationModel(tf.keras.Model):
     self.head = head
     self.mask_scoring_head = mask_scoring_head
 
-  def call(self, inputs: tf.Tensor, training: bool = None
+  def call(self, inputs: tf.Tensor, training: bool = None  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
            ) -> Dict[str, tf.Tensor]:
     backbone_features = self.backbone(inputs)
 

official/vision/ops/anchor.py

@@ -304,7 +304,7 @@ class RpnAnchorLabeler(AnchorLabeler):
     return (ignore_labels + positive_labels + negative_labels,
             positive_labels, negative_labels)
 
-  def label_anchors(
+  def label_anchors(  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
       self, anchor_boxes: Dict[str, tf.Tensor], gt_boxes: tf.Tensor,
       gt_labels: tf.Tensor
   ) -> Tuple[Dict[str, tf.Tensor], Dict[str, tf.Tensor]]:

official/vision/utils/object_detection/visualization_utils.py

@@ -878,7 +878,7 @@ def update_detection_state(step_outputs=None) -> Dict[str, Any]:
     else:
       detection_keypoints = None
 
-    if detection_keypoints:
+    if detection_keypoints is not None:
       state['detection_keypoints'] = tf.concat(detection_keypoints, axis=0)
 
     detection_masks = step_outputs['visualization'][1].get(

{tf_models_nightly-2.12.0.dev20230426.dist-info → tf_models_nightly-2.12.0.dev20230502.dist-info}/METADATA

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

{tf_models_nightly-2.12.0.dev20230426.dist-info → tf_models_nightly-2.12.0.dev20230502.dist-info}/RECORD

@@ -24,7 +24,7 @@ official/core/registry_test.py,sha256=8iQvUwOajmJ_ajHMjV7GkhhCX6ymNCTRcq1i7rQnmu
 official/core/savedmodel_checkpoint_manager.py,sha256=6q0IBaNsQ3IsxF08dAsUGcQoQ7Vo74ELJHH_znT6Dj0,9225
 official/core/savedmodel_checkpoint_manager_test.py,sha256=XGLlc_CimGMJLHL1-vtznBqEJ4j7S9IE7iS5Obo1lJg,4045
 official/core/task_factory.py,sha256=mFwJ3jHEj2FjWiEf3w22V5n2qwPywmXFeR9FZvYZLMM,2513
-official/core/test_utils.py,sha256=0uI2Ntypheyjqe_9NCLjOMv4F4_ozrDIHeyzlTWy2M4,1802
+official/core/test_utils.py,sha256=bsjidTNPtPOTGyDQRmi2MCgi1OjDw7PW9BmAYnCiQyM,1877
 official/core/tf_example_builder.py,sha256=eI34jkqX378mvtIVsSldtIlsBSfonqNwleh5pQrQmpc,4623
 official/core/tf_example_builder_test.py,sha256=1g9iB17VqsO36hyXbfkprPHUzggU6dF5mX1XsoLlXFQ,5793
 official/core/tf_example_feature_key.py,sha256=WJC1HAQ2y2h7goJ2Pbf2ejUYvi2CQ5jUGbRv7vJcQQ0,2095
@@ -223,17 +223,17 @@ official/modeling/multitask/interleaving_trainer_test.py,sha256=fw7wJRtghPMIuAOf
 official/modeling/multitask/multitask.py,sha256=pciaApxYL0X1Y6hj3wrSWHzT5Q-T7T__47vBNnYzhlM,5938
 official/modeling/multitask/task_sampler.py,sha256=unXwSKOqg1XoJYdF7c5Nxd0QFUtxjNSzOqhPJxBUp6Y,4887
 official/modeling/multitask/task_sampler_test.py,sha256=yUqlPXYAyH5JmyeVCPVJ2efQZfVKMO7y5m4KEeGp090,3027
-official/modeling/multitask/test_utils.py,sha256=iqjyetzUZRe4BWIiSUsAW9gjTFSXZxo9PWqNQFNfVRg,4155
+official/modeling/multitask/test_utils.py,sha256=CxvKg4YTP1B6WynXudcZrlrD-rXFj3zRGUz0bj2RGXE,4305
 official/modeling/multitask/train_lib.py,sha256=_5KBdUlUhotyiwWsdJnTm2e9MnR2UI0mCTAy1VRkgpI,10832
 official/modeling/multitask/train_lib_test.py,sha256=YGS0Aa1PLfH2FQtwm8v6PYoqeSSdusvllrnunUfREpk,4756
 official/modeling/optimization/__init__.py,sha256=mKS7Ujww2sec6_n_KXFktt5iJdAVIGHPTAaMYWdQa3c,1201
 official/modeling/optimization/adafactor_optimizer.py,sha256=Qnw1VSjRI5HJbsKBv3Q83tRhFno35Vfu0G1jNxA8Ofc,792
-official/modeling/optimization/ema_optimizer.py,sha256=HERzbA_EhUzLfWN92O6bAAjU4EWHTpQgjmcIqozR0zI,10511
+official/modeling/optimization/ema_optimizer.py,sha256=ecf1C_y_GcyVKr_MrEMlbjVUzkt9b3src8o8J-1qUwA,10508
 official/modeling/optimization/lamb.py,sha256=bANup6_NWFaITBMComKKCXpzU8_hBDcr3AARZiNP0m4,10137
 official/modeling/optimization/lamb_test.py,sha256=hTLunlnOr-VWzRUTRdkeDWKdSBMCmQJcrEH-gP75Fuo,5928
 official/modeling/optimization/lars.py,sha256=UO5YPPlaGR5lyrBH1FzUvVGzMefklioJXxGxWhHGWgk,7338
 official/modeling/optimization/legacy_adamw.py,sha256=8d0Zw4gDsDOub2vCU4t9rgFEZ7qVTUOdT82BGnq0SR8,5953
-official/modeling/optimization/lr_schedule.py,sha256=U-rFaMc2xT2NaPemy2UZvHpn0Seds5NRHwQKsd5TWD0,19097
+official/modeling/optimization/lr_schedule.py,sha256=skKKqQ51aEO5xRZWLrElqLUJEB2-Wd_6WLhA0wurPk8,18498
 official/modeling/optimization/lr_schedule_test.py,sha256=G_IV6daj1S5cKsypVaR8NfLbrB7DvtFZ4HQCtP0I-oY,3951
 official/modeling/optimization/optimizer_factory.py,sha256=8s1hMSgyQZKk-QoNI1WunOSUyTu5Sri2guLvsgl9Mqg,10475
 official/modeling/optimization/optimizer_factory_test.py,sha256=MiHFcKvsSPsgHnBxk8T_CwzakcGnRLvw3bJ7PO1FlLI,17015
@@ -301,7 +301,7 @@ official/nlp/modeling/__init__.py,sha256=bL8szAeBlvUmdY-yNRk-Nsurs7USPP3K4_pnyjj
 official/nlp/modeling/layers/__init__.py,sha256=0YkQXKDdGZl1H8K0pKb93ex3L6nY_vP4XFNXPKH8OJE,4796
 official/nlp/modeling/layers/attention.py,sha256=fJj_bsWgosN8p4Yy1eFToKaAuQICef_L_iwkHeg5gfw,3896
 official/nlp/modeling/layers/attention_test.py,sha256=vMBf1RRgx3tg3CXNyi6XW9nwG6KKnea8IzJ58udwN1o,3526
-official/nlp/modeling/layers/bigbird_attention.py,sha256=pHrAfII9B-BbSTcY3mx4B5s-BwRDanNEiIysc1XfIds,21026
+official/nlp/modeling/layers/bigbird_attention.py,sha256=XOcuap9IKvpXSkWtiAztqsutIhdJ66bylPKhDvzNA5A,21101
 official/nlp/modeling/layers/bigbird_attention_test.py,sha256=RVo4pNazC8ktRE7YI4HdyrM0X69pFmz9b8B-9lp7ix0,2196
 official/nlp/modeling/layers/block_diag_feedforward.py,sha256=KYzPAyZDCyA1jTL7P9gTNAJVRypngbkZ5Sv5QAbfTRA,7233
 official/nlp/modeling/layers/block_diag_feedforward_test.py,sha256=JrSs8t7Ry4ZZHYGrk_W4Tp0EBGqhboHRtKbj4jT68Lo,4171
@@ -333,11 +333,11 @@ official/nlp/modeling/layers/on_device_embedding.py,sha256=Z7sDSzdVe7CaMLfPe2DEP
 official/nlp/modeling/layers/on_device_embedding_test.py,sha256=t5s90xeiHB7WVJKb-YKw63hx2Pgjf9UcvPdZRpy_CwQ,8579
 official/nlp/modeling/layers/pack_optimization.py,sha256=b1XKyUtxeygyVzsDJveZIm-vgsdy-xznSRoXjyZ_Mzk,10279
 official/nlp/modeling/layers/pack_optimization_test.py,sha256=cuNoFFxwLB9cT5hw0lfy_aA7MTHOO3Mh231UJO65_Ak,2785
-official/nlp/modeling/layers/per_dim_scale_attention.py,sha256=GifLuuEymebNBcExgeuSvX8aqReGRBpNZNGSpO3AyJE,3331
+official/nlp/modeling/layers/per_dim_scale_attention.py,sha256=qQIuonWRU92kO0-RHqoI4BHEWIm_v-XnozXYWxUjmbM,3406
 official/nlp/modeling/layers/per_dim_scale_attention_test.py,sha256=7GZE3CyIHZx9x_s2Dwn5YNODLbyjhPH-SpAJojDET_0,1751
 official/nlp/modeling/layers/position_embedding.py,sha256=X7yaGV7KoHN-t51wk3jWRVg6w_gpgoYnhLTxUP-qEMs,11337
 official/nlp/modeling/layers/position_embedding_test.py,sha256=CueBbGQYpirkZpjlID2xAd4pmh49GWssx2QN7jKU6H8,8009
-official/nlp/modeling/layers/relative_attention.py,sha256=NaWtSHVUjb5CD8AZVYU8M3B8s5CaRWGQDjs6PgzqGfk,20472
+official/nlp/modeling/layers/relative_attention.py,sha256=flArMqRqCtHIQd6yWisscgclaEq6S_HxtpGV81tHrgw,20547
 official/nlp/modeling/layers/relative_attention_test.py,sha256=rTN71dKVytcIceiHRUPLPKdDogPnHJ_GzVW9bU_OXHQ,6685
 official/nlp/modeling/layers/reuse_attention.py,sha256=QjNS2hipbSKMSjqsldmNMjQ_3plYhGX_xj_X87_iQ08,25658
 official/nlp/modeling/layers/reuse_attention_test.py,sha256=340CNOwxcNxf2RKZnBdS4rS-rv0VkrbiCww5TSdjLWs,14319
@@ -348,14 +348,14 @@ official/nlp/modeling/layers/rezero_transformer_test.py,sha256=NmD3v5xcnpRxTe9dk
 official/nlp/modeling/layers/routing.py,sha256=BdNat8l8PWXE1eIdYXBh41Uy89BJa1b4OR7x9w69W34,4459
 official/nlp/modeling/layers/routing_test.py,sha256=iEx_88DCG7G62ML_TCGR1rS5CXlNSZEGZ8E3Z8UkkEI,2216
 official/nlp/modeling/layers/self_attention_mask.py,sha256=vFcKClhCma8R3DZ3V_f-5oIXr36eOFp4s7IPpXGyxWA,2163
-official/nlp/modeling/layers/spectral_normalization.py,sha256=iVM9Yn3RYDV5NdYxQgZswu74BK5kTkWwl9X3r92kAtA,10578
+official/nlp/modeling/layers/spectral_normalization.py,sha256=2y-Ay8MkDzl3DytPInMqUjx_ww7UfmYgOCvRZjjAYxA,10728
 official/nlp/modeling/layers/spectral_normalization_test.py,sha256=acyfgLtuzKEpws14u-aqWqni6D41Z1gdq4hRfmvBQwM,3111
 official/nlp/modeling/layers/talking_heads_attention.py,sha256=shkxv2FzH76jSzTdxtZyPgu8kBnuNzNZlaWGti_8xwY,6903
 official/nlp/modeling/layers/talking_heads_attention_test.py,sha256=oqgiQPsNLpckHJy1BYN56O5P1jXzHyYj9Y2JaejVqzU,7012
 official/nlp/modeling/layers/text_layers.py,sha256=soVaf1cqb7s25Fk0i01wCJE97sscY-PFfVNjpcuJjhs,32516
 official/nlp/modeling/layers/text_layers_test.py,sha256=0aP-lzjnuWrn26HuWncnj7ASIEMFRVrR6W_Mwob3WdE,24245
 official/nlp/modeling/layers/tn_expand_condense.py,sha256=yzCgGwXfTC0Mumn8HmbczZ1pG9JYkpIE2LLb702ADUU,6703
-official/nlp/modeling/layers/tn_expand_condense_test.py,sha256=5yyhG_K_S7J4frq5sq7NmEfjwwiHvVocK-GZB0vqfeA,6434
+official/nlp/modeling/layers/tn_expand_condense_test.py,sha256=LXu-Of1ck-SeTCfUo9gGfL_CFTxPLVpJTQw8NtCX0TU,5890
 official/nlp/modeling/layers/tn_transformer_expand_condense.py,sha256=RK5vryB7W8cU58_BWzfee2id1RB6BT2NVhP9FlYkXYI,11022
 official/nlp/modeling/layers/tn_transformer_test.py,sha256=frCu6LKapVWXw0YLs8_7KMp5qxEOHNBy7yJIZiE17qc,8883
 official/nlp/modeling/layers/transformer.py,sha256=1URI5Z3ZV6DQ2oTgdJIc6EXbgutSlDZ7l-gYoV1HCcE,19414
@@ -373,7 +373,7 @@ official/nlp/modeling/losses/weighted_sparse_categorical_crossentropy_test.py,sh
 official/nlp/modeling/models/__init__.py,sha256=bOt0Hr5hjT9K7NQK-hzG9fsye3flx3U4X8RhkF8neN0,1654
 official/nlp/modeling/models/bert_classifier.py,sha256=qR8trNtfxhO_FWKUyTR43q3_Xu3LNWMxRUlvN3uQxhk,5896
 official/nlp/modeling/models/bert_classifier_test.py,sha256=uAXgZgK2Zq2dtIsbDJWqUO1Aro0Beyjc4rFMBLE37OA,4815
-official/nlp/modeling/models/bert_pretrainer.py,sha256=Is7T32hB5D7Klm1u636dRbFpZh6SX-mzgbe9cpkisAg,11405
+official/nlp/modeling/models/bert_pretrainer.py,sha256=upQUW4JBvFPqibsSiMxyXL30Qwn8GGaZsR4ekHgi7ZU,11480
 official/nlp/modeling/models/bert_pretrainer_test.py,sha256=g4yaAdaoNGcfzt2vDMhwJleTRH9pLNzCYOGol5o7D84,9900
 official/nlp/modeling/models/bert_span_labeler.py,sha256=J7p_RZc_B89UUZbJ8Z3vrbujIwJWuMZ8_20Ti-4bsHI,4997
 official/nlp/modeling/models/bert_span_labeler_test.py,sha256=s75mgXjAcZrOTnsOUVkcxuAS0v00nomOlnvf3EW80UQ,4769
@@ -381,13 +381,13 @@ official/nlp/modeling/models/bert_token_classifier.py,sha256=IoaQsv_uFCnmlfxTRFU
 official/nlp/modeling/models/bert_token_classifier_test.py,sha256=bkL644Hp57GlW_NcQlrPobF2YO_ueJLLLauuxOsnLmU,4835
 official/nlp/modeling/models/dual_encoder.py,sha256=_yUt8Ox4iWhJCF3zqIbBDba0WxcAvd1PrVjQACJZRjs,6583
 official/nlp/modeling/models/dual_encoder_test.py,sha256=XZ1AGzR55WFWmo0VstV4eWcvkBzJlHyl_6H9nQIM7Xk,5120
-official/nlp/modeling/models/electra_pretrainer.py,sha256=Qv2kwobpon8iZUq6bFHrN-ck_SeHiSeVhBqryQA7oUg,12961
+official/nlp/modeling/models/electra_pretrainer.py,sha256=Eg9aOziKn761EDEUs3gqxUOgdq1nHArjUX_l2Vcc2VE,13036
 official/nlp/modeling/models/electra_pretrainer_test.py,sha256=ArnVb5gURE5xj88DIq6HsZib0bcFk46ePpsd9Nqa6Fo,6503
-official/nlp/modeling/models/seq2seq_transformer.py,sha256=FkH8yk3zfqVeNrKtL__5Ow23w0dkgzCGT5Ipghjn7oc,25716
+official/nlp/modeling/models/seq2seq_transformer.py,sha256=4oEIFMEWavyDRRvtaz2g-bAKiD3ehdrLM6MSzNfOw90,25791
 official/nlp/modeling/models/seq2seq_transformer_test.py,sha256=3fGfPITjC_xdNL8B8U5_qVPoYXS6_7Ic3jlBKaeebrI,5399
 official/nlp/modeling/models/t5.py,sha256=pQk5Gleeq3Sr90tVQQ5uZlftyViVNhqY2l6d3f0DUYc,57312
 official/nlp/modeling/models/t5_test.py,sha256=UsF6AMuTRRf_OX9xAcbwIRsuYukp8KS-qtVHAy4Y2jM,27024
-official/nlp/modeling/models/xlnet.py,sha256=OvTkl6xGYDn-IjxZkTP0p-t5mZVgfysBzP-k44dUji8,11779
+official/nlp/modeling/models/xlnet.py,sha256=kJkmNsKavBKlPPwiklbGA0okakCJYky3M87WUjtVQrc,12004
 official/nlp/modeling/models/xlnet_test.py,sha256=VSG37nAOuZkgurNlNDCAXXpOZ-KeEXm8YRQvBWS1UzI,12608
 official/nlp/modeling/networks/__init__.py,sha256=TQT6o5ZZUtc_ofg8gSUq30bA4FtId6Mz-z39dEF73Rg,1839
 official/nlp/modeling/networks/albert_encoder.py,sha256=LgrfJB5jYZsNQ2LPpZhoeko3TekCwQuT_taPVxjhLa4,8877
@@ -477,7 +477,7 @@ official/projects/centernet/losses/__init__.py,sha256=1ToRMjre4mErL4Ek4_dMVxMjXN
 official/projects/centernet/losses/centernet_losses.py,sha256=rN7tyZ0VwNMgCFUKn7ce9HCAPBcFhxQIVug5jPgPRXM,4548
 official/projects/centernet/losses/centernet_losses_test.py,sha256=VlUVFqFxHhyMIjECHepmzsvePQvRpMG9f6Qa-loyFmM,3777
 official/projects/centernet/modeling/__init__.py,sha256=1ToRMjre4mErL4Ek4_dMVxMjXNPossNXggV8fqbISao,609
-official/projects/centernet/modeling/centernet_model.py,sha256=9DEPNcZSLSjcuA7ev-YjVQ5TP9yKdGzU3G-90lXhrBk,2541
+official/projects/centernet/modeling/centernet_model.py,sha256=uAXOCcotwzXKQvCbXHCcGPPQA2OlyMq56fdJ7VqWEHE,2616
 official/projects/centernet/modeling/centernet_model_test.py,sha256=WWUNxWXcu-zpI2ZcdeXrfOqcxvEzeODvRn5lnPCt6mc,2575
 official/projects/centernet/modeling/backbones/__init__.py,sha256=1ToRMjre4mErL4Ek4_dMVxMjXNPossNXggV8fqbISao,609
 official/projects/centernet/modeling/backbones/hourglass.py,sha256=M5aOsOHFGcHaTTAdEyRv5s0ThTqDh-hYAOs91KT6490,10285
@@ -517,7 +517,7 @@ official/projects/deepmac_maskrcnn/modeling/__init__.py,sha256=1ToRMjre4mErL4Ek4
 official/projects/deepmac_maskrcnn/modeling/maskrcnn_model.py,sha256=_hRwbPrjII6WhhrK6Ehu1Zpo8XBmQQf_Evzmpy0VflY,9656
 official/projects/deepmac_maskrcnn/modeling/maskrcnn_model_test.py,sha256=0_zFQo30yEDME7o_lo_KJqfegmIJjlxUl0Y-saHdqIs,5731
 official/projects/deepmac_maskrcnn/modeling/heads/__init__.py,sha256=1ToRMjre4mErL4Ek4_dMVxMjXNPossNXggV8fqbISao,609
-official/projects/deepmac_maskrcnn/modeling/heads/hourglass_network.py,sha256=vRWvh9WmxadsjbJwn6TF93R7fdEt3Mhsa7doXCFW-rA,21630
+official/projects/deepmac_maskrcnn/modeling/heads/hourglass_network.py,sha256=dochwOGv2S_O_Fd8ew0_nTn58wNh_Q6vLWKb9cychyg,21705
 official/projects/deepmac_maskrcnn/modeling/heads/instance_heads.py,sha256=fsCbLANUxJ5QRVUVhThZgymJIWLQC39Zr0uooP7GO2g,11466
 official/projects/deepmac_maskrcnn/modeling/heads/instance_heads_test.py,sha256=M8iuImwpTHnnVXDLNCHVVk7NunIzKJrMrRyQc9_RTtU,3172
 official/projects/deepmac_maskrcnn/serving/__init__.py,sha256=1ToRMjre4mErL4Ek4_dMVxMjXNPossNXggV8fqbISao,609
@@ -651,7 +651,7 @@ official/projects/qat/vision/tasks/retinanet_test.py,sha256=Kl8s3-ag5ME58QthtwiY
 official/projects/qat/vision/tasks/semantic_segmentation.py,sha256=mhQ7hb3M4e5f71kLDa82aWyqQA-gUIfOw9GFUSoD05o,1552
 official/projects/roformer/__init__.py,sha256=wrNoNRP6wGQBFDLVCYB01k1HgwOGlQ8iZnfqimxlDyk,610
 official/projects/roformer/roformer.py,sha256=6FV81XLin-7TqZQLwmS2gdOpfb4S8jSXvoo1Og3vZ18,2002
-official/projects/roformer/roformer_attention.py,sha256=FwtmYS8xuGKBiVLtWHzP3R-rkCMbXs3R_-rey3aN_HI,4563
+official/projects/roformer/roformer_attention.py,sha256=cR1YY-dWLGnFN4yyof9kYH9v_DM09a8SCA5THdjK52U,4638
 official/projects/roformer/roformer_attention_test.py,sha256=IQpM6PWCRCilVBxpAk80tVnqhwcUVM_t1lio871yBkQ,4037
 official/projects/roformer/roformer_encoder.py,sha256=InBAjuLgz_bAH5Sgqe32cgdhmh6bSsLMD-NaqFYqvwE,11374
 official/projects/roformer/roformer_encoder_block.py,sha256=Cfe3NStzPrHK8JsL3DjoM7jemGCY2aWGu6a6d1Z1r9M,13507
@@ -662,7 +662,7 @@ official/projects/roformer/train.py,sha256=zDMF_2jbLB8XrOG1ea_sGIYSS0q-q8k3QVkON
 official/projects/teams/__init__.py,sha256=1ToRMjre4mErL4Ek4_dMVxMjXNPossNXggV8fqbISao,609
 official/projects/teams/teams.py,sha256=JwhdzJnbsZkMHefM0vBP4eAPadh11LbMkEpH-UFC-po,4010
 official/projects/teams/teams_experiments.py,sha256=gLFMx-Osj5LqgcuBAwOJkWOM4dG3Z0UgOHARMB42y5g,4424
-official/projects/teams/teams_pretrainer.py,sha256=73-6zqeqs_QDLSj6xAqCieheI1yfYxWSt1I32I1X16A,18608
+official/projects/teams/teams_pretrainer.py,sha256=MZv7Tc2WHPPoWcffI3Kk9l_UbdnkBJkzrdbFPKxeyag,18683
 official/projects/teams/teams_pretrainer_test.py,sha256=IwABozbqvlDfYujZiGIIIaoX4Fi4E4GN43PXFxTpRPU,7436
 official/projects/teams/teams_task.py,sha256=VmNuUhOkyttFImSh85CuqdnGg7og6wrx12M-vJQcEfA,10170
 official/projects/teams/teams_task_test.py,sha256=_3A-z-zLhH8Y8LdPbBaQJRbzAJiune-8JpZA94Fvqow,2184
@@ -761,8 +761,8 @@ official/projects/yolo/optimization/configs/__init__.py,sha256=1ToRMjre4mErL4Ek4
 official/projects/yolo/optimization/configs/optimization_config.py,sha256=NmxuFW64_MYFcBF_SxN-34rkQ7PxAswpxpufBL2GIPA,1998
 official/projects/yolo/optimization/configs/optimizer_config.py,sha256=0YZgwyEQKfEwP44x9wzZ7N01I0bdawqEOgWmezora0Y,2283
 official/projects/yolo/serving/__init__.py,sha256=1ToRMjre4mErL4Ek4_dMVxMjXNPossNXggV8fqbISao,609
-official/projects/yolo/serving/export_module_factory.py,sha256=O_xdB7h7ra_mpKXGEaXW09iiqhh2LiR2B25Wj1j-wTk,9599
-official/projects/yolo/serving/export_saved_model.py,sha256=3NZbcxfkWzwpAEkLs5vOven2cxy3nnYix2mPnIp-05o,4368
+official/projects/yolo/serving/export_module_factory.py,sha256=tRtJbbaLxV1BLm3OTTkyzkkyRrGxOTJdcg4gCuInW6k,9910
+official/projects/yolo/serving/export_saved_model.py,sha256=gOxopjUl4Axp_Z1ArE_QeUW2qM0JiX71I040lQOKDrY,4285
 official/projects/yolo/serving/export_tflite.py,sha256=Pt1rkKu9TxVYB8HtLadJQ_buWzNXqWL8Zp0TFDPUXzw,902
 official/projects/yolo/serving/model_fn.py,sha256=J1i-UPX1TuPqPaBWu2HyhtVHthq6VRbsnR98Sjrs0iY,3008
 official/projects/yolo/tasks/__init__.py,sha256=1ToRMjre4mErL4Ek4_dMVxMjXNPossNXggV8fqbISao,609
@@ -830,13 +830,13 @@ official/utils/misc/model_helpers.py,sha256=zqgwrD_oi23qBaY7PnMIBKWHqUsNGJRwgral
 official/utils/misc/model_helpers_test.py,sha256=WS0Eb0pPGk746JSDXRJyZ0Yu95_RvwwmN0uJakRf-ls,4549
 official/utils/testing/__init__.py,sha256=1ToRMjre4mErL4Ek4_dMVxMjXNPossNXggV8fqbISao,609
 official/utils/testing/integration.py,sha256=Wac7nJTxnaK-XHt5pNbFD9eaP4CwZ9mnr7ZOMkJosEw,2220
-official/utils/testing/mock_task.py,sha256=NibyG1n5N_FisqTFwFqarGetsnpMH-dQ5lRQzD3kJps,3227
+official/utils/testing/mock_task.py,sha256=HPDArX_WR_ZZqom2iCSwn_zy8tbSjOGIN38KLL36bhw,3302
 official/vision/__init__.py,sha256=g2WviTDdhFjzzfQROU7MV-xdCchqoyXvs_RlIE8fiTc,744
 official/vision/registry_imports.py,sha256=1T4yHR55Tct02_E5P71jR0LtAkft84IZC2Zt2hwm7dI,760
 official/vision/train.py,sha256=QQlYD8NtC3_JYLvPhV3i0GoqTUOaYoEHRHi2UoNracQ,3724
 official/vision/train_spatial_partitioning.py,sha256=sXVjCU8r29aXP6vc8kp8DxciEhfdn3APoyQOr-iCQZ8,5725
 official/vision/configs/__init__.py,sha256=5Uk9jeIYHTEHXs8kBrq-vMa7cgg_JpLgSYuVaLPyUt8,1045
-official/vision/configs/backbones.py,sha256=gHHldpb1pv7nqgBIQPZwV_bCcfSOquv_rfc8_HEIgUg,5070
+official/vision/configs/backbones.py,sha256=E87hCWSPU07fqS-SPkUYsFeapkSIjGQBIAU_dq-Knw4,5197
 official/vision/configs/backbones_3d.py,sha256=YGmeo0y1sedOQpfXdcXUubmgPpDS8IlyjM1yBSBTxAU,3650
 official/vision/configs/common.py,sha256=S_SYgZs9mSU5udq9xUuGlt52dzhdemobVzBFrGfhNPc,4621
 official/vision/configs/decoders.py,sha256=eSUlND3gFJK_z9yrHdoRzJN-PY_8v2_TcgwY36G9OqU,1948
@@ -913,11 +913,11 @@ official/vision/modeling/classification_model_test.py,sha256=M8t4MA5cI42rmfTqGsX
 official/vision/modeling/factory.py,sha256=EfuDNq9629SIRYn01Dy1h4gGyjPa2wfPROhQTL-9dQs,17437
 official/vision/modeling/factory_3d.py,sha256=Gtz4xWxPQ4rmMn360lTMgw_Xjha5q6UYawotcerXpYM,3530
 official/vision/modeling/factory_test.py,sha256=I1aevMcMbD6dihHfU5ohFV4R2lHEdVLvdVra7QE-kU4,5286
-official/vision/modeling/maskrcnn_model.py,sha256=q02E6boSm15k5xJl_ggmMXaXebBOREbAccqolAohAZQ,20658
+official/vision/modeling/maskrcnn_model.py,sha256=1eb4mIbZeYA1ofBk2MosF6VsNRDJph9gBPCxU0Zb95w,20733
 official/vision/modeling/maskrcnn_model_test.py,sha256=eJ0WtCGFztZj7zbgAd0VDF4_PzsnLgPK0k4r1n85X5k,14834
 official/vision/modeling/retinanet_model.py,sha256=kT9j8khuYxZeRbJ070Ys84y1hU412VzZcHpJVpBEe4Q,9695
 official/vision/modeling/retinanet_model_test.py,sha256=_PTpYedMG0L8rq6Xydh1qR0vKpvE4uDu1M6HZF4M3Ig,11148
-official/vision/modeling/segmentation_model.py,sha256=QMhyv8CW9CGnpMLrCvPOFdQJBEWIt0EvLpsdONxX0h0,3353
+official/vision/modeling/segmentation_model.py,sha256=Vyhq0G2_Smf0RCu1gRhVhu3C145r1lRRWAie6keU60M,3428
 official/vision/modeling/segmentation_model_test.py,sha256=Z5toOhb3dIyU7DE38pLOalwBYKx2n1prIxawkOu4GHI,2807
 official/vision/modeling/video_classification_model.py,sha256=em4Yt3grA7Y7GA9bSzPSCCrc1_l_RNUu6qEWmCZUrZ4,4703
 official/vision/modeling/video_classification_model_test.py,sha256=eX0sBnEdLwPV7qPfR9nSApRw1qd0iUPew1yCLMn4j4E,3261
@@ -942,7 +942,7 @@ official/vision/modeling/backbones/spinenet.py,sha256=uIaJdINdVnvyh5xIs_1oRMajQA
 official/vision/modeling/backbones/spinenet_mobile.py,sha256=vVpqM7oKSQMsCTJtZZBPWnbKei5oQHq33woOwIHiF8w,20774
 official/vision/modeling/backbones/spinenet_mobile_test.py,sha256=UvLU7Jutk7f5q8Vz3jflWHlC-itRMvNJtB5-cXA0JxA,3948
 official/vision/modeling/backbones/spinenet_test.py,sha256=YhQRhc3bDhpNd-EUQYxDkLSfrOhKjcpeUeh-nQ__pZg,4724
-official/vision/modeling/backbones/vit.py,sha256=BUetPmiFKWC8C119h2BM9ROUyAH57RNTpnd_6ZqM3QM,14162
+official/vision/modeling/backbones/vit.py,sha256=eO18uyTYZ7Uy9FKp-m1zVpNqy4AHwIMBvgJqpwJ6IiY,14428
 official/vision/modeling/backbones/vit_specs.py,sha256=Qfa5Oecd4rCWA73m8X_g8J_w1fpCZsNRSUw0qLIhO9Q,2412
 official/vision/modeling/backbones/vit_test.py,sha256=aP11sEVNfOvb8i3VfRd3k_3qHWqFvunEYO9DRtCTZrA,3453
 official/vision/modeling/decoders/__init__.py,sha256=F7uTs9fJjPKacKgZnuRRm21t1N4cV_c4o76dDQS9VkM,815
@@ -959,7 +959,7 @@ official/vision/modeling/heads/dense_prediction_heads.py,sha256=Xmtb4moAiinD3-wC
 official/vision/modeling/heads/dense_prediction_heads_test.py,sha256=8O1K4iksJ7IBRVfiHPjaolcAYd3sl0DxOoTljfJf_c0,7227
 official/vision/modeling/heads/instance_heads.py,sha256=P580eVgYex6ddcua1zTkXEmgOqPd1xnmqJbrM6Dz6Hk,17819
 official/vision/modeling/heads/instance_heads_test.py,sha256=Hjsb196FnfJtG4NMXmZuibo_sgBINPy-qzPyXZ_Ru9A,4197
-official/vision/modeling/heads/segmentation_heads.py,sha256=H7CMqC1i28-IDCovFCYbKbwo1sKyNVlLSqbBK5OHU_k,20099
+official/vision/modeling/heads/segmentation_heads.py,sha256=ABDjiLFo7pmyV4uodPRXQrNxWHhzhcQ7xHgrjEsmPJk,20174
 official/vision/modeling/heads/segmentation_heads_test.py,sha256=VJOg0_VcpdyMj3HYfxG1BcpHm0D9B5yEw6kZISiZCY8,3802
 official/vision/modeling/layers/__init__.py,sha256=Yk0zkJVVgwi--tv5-NKIAKHCzpFw9tKrg5wYm5CZwh4,2597
 official/vision/modeling/layers/box_sampler.py,sha256=hsFoo9lsD57LDUnj0SlIoM9QNg26yqQxvADPtQRVhV8,3401
@@ -970,11 +970,11 @@ official/vision/modeling/layers/detection_generator_test.py,sha256=-XXoY6dyjE1R3
 official/vision/modeling/layers/edgetpu.py,sha256=AmfMIxT1XCSZo0zXoZO0h-WOdxNYYYZo4lGf_d7DL6g,16727
 official/vision/modeling/layers/edgetpu_test.py,sha256=aIOtLXD1jCExn3tVFS-q6yiCZV7yIxCvG06p7fQdCgs,10165
 official/vision/modeling/layers/mask_sampler.py,sha256=1INadZeRmqZoiFeEHiE_IMq39k8f2S66VvoyMxcM1mk,7914
-official/vision/modeling/layers/nn_blocks.py,sha256=GhBsCH0jYwi9fyBKh4S_X4WyU42k-6ODWo7k7oxrIiQ,71345
+official/vision/modeling/layers/nn_blocks.py,sha256=wU0OK-l1TPeMpnNmJwmUxypNpag2wXsaIhYp2NunmLo,72079
 official/vision/modeling/layers/nn_blocks_3d.py,sha256=b0h_wWmzgrRUaG03eziHoAqBpqZJyfSn8riHfFbEKzo,10565
 official/vision/modeling/layers/nn_blocks_3d_test.py,sha256=pSmDiikoPq83gZIP9hyZ0UbmJxMxVxippuLd38mHR74,2028
 official/vision/modeling/layers/nn_blocks_test.py,sha256=-NNR8zVnkgkyLTyNzOL2BnkvH5IWz1ggO3sXCeYe-E4,33825
-official/vision/modeling/layers/nn_layers.py,sha256=OkH9xwhgVQIlJbnXNrCL3PBBRLJ9Xz26VJcukhZ2KaI,52327
+official/vision/modeling/layers/nn_layers.py,sha256=Bv18J7Sa-FHLrlIru9vSXiJDOlCw2M9xtoOQH0_uuPM,52972
 official/vision/modeling/layers/nn_layers_test.py,sha256=NMthj6XogpbTcbhjxsMiO8WEgq6ha4uLJwwqMDOEY5I,13309
 official/vision/modeling/layers/roi_aligner.py,sha256=v09HensTcogJDNzw2dIUtvbQhoPzaxvSVLVsbISwYrw,2533
 official/vision/modeling/layers/roi_aligner_test.py,sha256=p5_zu8nXS2L1mV5zX9iGIgC7OQc1gd0Fw7o1AffyzhA,1275
@@ -982,7 +982,7 @@ official/vision/modeling/layers/roi_generator.py,sha256=VGNMrm0_td1gaE7ENK-XzGEy
 official/vision/modeling/layers/roi_sampler.py,sha256=3RZBvQ3esodqT_BzDvAG0r4x7mxOHKoxLwSk0V_JNEs,9996
 official/vision/modeling/models/__init__.py,sha256=dKnmOzbcEAJcah2R85pfGE37VZYY090oMy9ol4eHk8k,1020
 official/vision/ops/__init__.py,sha256=1ToRMjre4mErL4Ek4_dMVxMjXNPossNXggV8fqbISao,609
-official/vision/ops/anchor.py,sha256=jW8F1js-bFCo4JGP4rFw4ti6mHhUuEfTEkQ1mvLhfe8,19027
+official/vision/ops/anchor.py,sha256=iS97q20CKozhg9tpqFkkDGZ66MqrIfCYjSk3ZjndrRM,19102
 official/vision/ops/anchor_generator.py,sha256=WwASWaE5uR9oIB69lnK8mW2xUfRnLY0LsgyVX1BwsQQ,7234
 official/vision/ops/anchor_generator_test.py,sha256=aihKzRV8SQFhNSzS9_SpJNSadL4KcUcj0vT5dgmnkIQ,5286
 official/vision/ops/anchor_test.py,sha256=UHue1kHD184Lh5eJzWQ2w0UAZk9FjolYJ9Eg6b1hIbo,7623
@@ -1020,7 +1020,6 @@ official/vision/serving/export_tfhub.py,sha256=ES-BDGUjhfv8u93S1RKnY-ZJw-g25NAxK
 official/vision/serving/export_tfhub_lib.py,sha256=x_wtofpZH-eb8y92IBdPexjta0ecG6l-9HoHOePzt0g,2870
 official/vision/serving/export_tflite.py,sha256=a-j5An1DigxPxpIg8IcqH2uQzaSXIQRL9irD2HmkAVg,5105
 official/vision/serving/export_tflite_lib.py,sha256=3zOA1W1-yLcHCzngtgauyyiyidOPZXuJR35HSDP3Rwo,6907
-official/vision/serving/export_tflite_lib_test.py,sha256=2TW-215vjJdLT6lPYUtS_tSxYYNhd0-JbfvPK46oUTc,7185
 official/vision/serving/export_utils.py,sha256=tPzbbJmbu7rLjiu5GRdCcaU6ADybNbAPH4DalJNURbo,4871
 official/vision/serving/image_classification.py,sha256=YWHm-VPnIpv4Hn973jM39saR2NZGB_4n8CSmJA9tIvw,2858
 official/vision/serving/image_classification_test.py,sha256=VwQc5WBRAMaB7Tm0TT2SJ2svoT2V9-OAWbZgT4QdsYA,4612
@@ -1050,7 +1049,7 @@ official/vision/utils/object_detection/preprocessor.py,sha256=Fklxck111TQEKsrWSM
 official/vision/utils/object_detection/region_similarity_calculator.py,sha256=yqQLLRT80IdAu3K_fILli_u1aL37lv0FpDtFcyRrPzs,4544
 official/vision/utils/object_detection/shape_utils.py,sha256=p3Q7e9gTTQNv1gnMAkuTkfXc6DYVB8mjk_Vjlq7bRlg,3608
 official/vision/utils/object_detection/target_assigner.py,sha256=fTkjedzhp_-RTUGR27tWbmnnLpW7F3lVCkZlr7Nv_9o,24198
-official/vision/utils/object_detection/visualization_utils.py,sha256=sZre2InIcPNO4jbtPjXw_ACwbT-wy84Ql57MmqMFX-k,34515
+official/vision/utils/object_detection/visualization_utils.py,sha256=ldXEkcW0O9Tdm-0HrRhtCi-gm7bho--2ky96DqfVH_o,34527
 orbit/__init__.py,sha256=aQRo8zqIQ0Dw4JQReZeiB6MmuJLvvw4DbYHYti5AGys,1117
 orbit/controller.py,sha256=RFcHfK4ZfgnscJC1uAlHav5SkehM2sFOeatyIpNnfTo,23667
 orbit/controller_test.py,sha256=sx1ok0Epn9gKMtN2bWCavAtIe5Q4WD_hLN2HqD-lsxo,29924
@@ -1084,9 +1083,9 @@ tensorflow_models/__init__.py,sha256=Ciz_YBke6teb6y42QyQTUBDdXJAiV7Qdu1zOoZvYiKw
 tensorflow_models/tensorflow_models_test.py,sha256=Kz2y4V-rtBhZFFfKD2soCq52hviSfJVV1L2ztqS-9oM,1385
 tensorflow_models/nlp/__init__.py,sha256=3dULDpUBpDi9vljpXadq6oJrWH4y6z42Bz2d3hopYZw,807
 tensorflow_models/vision/__init__.py,sha256=4y77XkHaH8qLls3-6ta4tMp3Xj8CLbB0ihH91HsQ9z4,833
-tf_models_nightly-2.12.0.dev20230426.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.12.0.dev20230426.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.12.0.dev20230426.dist-info/METADATA,sha256=WCI1jX2cl3TRpemQp8RhhfKVmRL-l858h8FArRTcKH0,1393
-tf_models_nightly-2.12.0.dev20230426.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.12.0.dev20230426.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.12.0.dev20230426.dist-info/RECORD,,
+tf_models_nightly-2.12.0.dev20230502.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.12.0.dev20230502.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.12.0.dev20230502.dist-info/METADATA,sha256=XiV5z3VirmxIf1tN16rxMxsHSF5dNi941h2yEcWKaYo,1393
+tf_models_nightly-2.12.0.dev20230502.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.12.0.dev20230502.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.12.0.dev20230502.dist-info/RECORD,,

official/vision/serving/export_tflite_lib_test.py

@@ -1,192 +0,0 @@
-# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for export_tflite_lib."""
-import os
-
-from absl.testing import parameterized
-import tensorflow as tf
-
-from tensorflow.python.distribute import combinations
-from official.core import exp_factory
-from official.vision import registry_imports  # pylint: disable=unused-import
-from official.vision.dataloaders import tfexample_utils
-from official.vision.serving import detection as detection_serving
-from official.vision.serving import export_tflite_lib
-from official.vision.serving import image_classification as image_classification_serving
-from official.vision.serving import semantic_segmentation as semantic_segmentation_serving
-
-
-class ExportTfliteLibTest(tf.test.TestCase, parameterized.TestCase):
-
-  def setUp(self):
-    super().setUp()
-    # Create test data for image classification.
-    self.test_tfrecord_file_cls = os.path.join(self.get_temp_dir(),
-                                               'cls_test.tfrecord')
-    example = tf.train.Example.FromString(
-        tfexample_utils.create_classification_example(
-            image_height=224, image_width=224))
-    self._create_test_tfrecord(
-        tfrecord_file=self.test_tfrecord_file_cls,
-        example=example,
-        num_samples=10)
-
-    # Create test data for object detection.
-    self.test_tfrecord_file_det = os.path.join(self.get_temp_dir(),
-                                               'det_test.tfrecord')
-    example = tfexample_utils.create_detection_test_example(
-        image_height=128, image_width=128, image_channel=3, num_instances=10)
-    self._create_test_tfrecord(
-        tfrecord_file=self.test_tfrecord_file_det,
-        example=example,
-        num_samples=10)
-
-    # Create test data for semantic segmentation.
-    self.test_tfrecord_file_seg = os.path.join(self.get_temp_dir(),
-                                               'seg_test.tfrecord')
-    example = tfexample_utils.create_segmentation_test_example(
-        image_height=512, image_width=512, image_channel=3)
-    self._create_test_tfrecord(
-        tfrecord_file=self.test_tfrecord_file_seg,
-        example=example,
-        num_samples=10)
-
-  def _create_test_tfrecord(self, tfrecord_file, example, num_samples):
-    examples = [example] * num_samples
-    tfexample_utils.dump_to_tfrecord(
-        record_file=tfrecord_file, tf_examples=examples)
-
-  def _export_from_module(self, module, input_type, saved_model_dir):
-    signatures = module.get_inference_signatures(
-        {input_type: 'serving_default'})
-    tf.saved_model.save(module, saved_model_dir, signatures=signatures)
-
-  @combinations.generate(
-      combinations.combine(
-          experiment=['mobilenet_imagenet'],
-          quant_type=[
-              None,
-              'default',
-              'fp16',
-              'int8_fallback',
-              'int8_full',
-              'int8_full_fp32_io',
-              'int8_full_int8_io',
-          ]))
-  def test_export_tflite_image_classification(self, experiment, quant_type):
-
-    params = exp_factory.get_exp_config(experiment)
-    params.task.validation_data.input_path = self.test_tfrecord_file_cls
-    params.task.train_data.input_path = self.test_tfrecord_file_cls
-    params.task.train_data.shuffle_buffer_size = 10
-    temp_dir = self.get_temp_dir()
-    module = image_classification_serving.ClassificationModule(
-        params=params,
-        batch_size=1,
-        input_image_size=[224, 224],
-        input_type='tflite')
-    self._export_from_module(
-        module=module,
-        input_type='tflite',
-        saved_model_dir=os.path.join(temp_dir, 'saved_model'))
-
-    tflite_model = export_tflite_lib.convert_tflite_model(
-        saved_model_dir=os.path.join(temp_dir, 'saved_model'),
-        quant_type=quant_type,
-        params=params,
-        calibration_steps=5)
-
-    self.assertIsInstance(tflite_model, bytes)
-
-  @combinations.generate(
-      combinations.combine(
-          experiment=['retinanet_mobile_coco'],
-          quant_type=[
-              None,
-              'default',
-              'fp16',
-              'int8_fallback',
-              'int8_full',
-              'int8_full_fp32_io',
-              'int8_full_int8_io',
-          ]))
-  def test_export_tflite_detection(self, experiment, quant_type):
-
-    params = exp_factory.get_exp_config(experiment)
-    params.task.validation_data.input_path = self.test_tfrecord_file_det
-    params.task.train_data.input_path = self.test_tfrecord_file_det
-    params.task.model.num_classes = 2
-    params.task.model.backbone.spinenet_mobile.model_id = '49XS'
-    params.task.model.input_size = [128, 128, 3]
-    params.task.model.detection_generator.nms_version = 'v1'
-    params.task.train_data.shuffle_buffer_size = 5
-    temp_dir = self.get_temp_dir()
-    module = detection_serving.DetectionModule(
-        params=params,
-        batch_size=1,
-        input_image_size=[128, 128],
-        input_type='tflite')
-    self._export_from_module(
-        module=module,
-        input_type='tflite',
-        saved_model_dir=os.path.join(temp_dir, 'saved_model'))
-
-    tflite_model = export_tflite_lib.convert_tflite_model(
-        saved_model_dir=os.path.join(temp_dir, 'saved_model'),
-        quant_type=quant_type,
-        params=params,
-        calibration_steps=1)
-
-    self.assertIsInstance(tflite_model, bytes)
-
-  @combinations.generate(
-      combinations.combine(
-          experiment=['mnv2_deeplabv3_pascal'],
-          quant_type=[
-              None,
-              'default',
-              'fp16',
-              'int8_fallback',
-              'int8_full',
-              'int8_full_fp32_io',
-              'int8_full_int8_io',
-          ]))
-  def test_export_tflite_semantic_segmentation(self, experiment, quant_type):
-
-    params = exp_factory.get_exp_config(experiment)
-    params.task.validation_data.input_path = self.test_tfrecord_file_seg
-    params.task.train_data.input_path = self.test_tfrecord_file_seg
-    params.task.train_data.shuffle_buffer_size = 10
-    temp_dir = self.get_temp_dir()
-    module = semantic_segmentation_serving.SegmentationModule(
-        params=params,
-        batch_size=1,
-        input_image_size=[512, 512],
-        input_type='tflite')
-    self._export_from_module(
-        module=module,
-        input_type='tflite',
-        saved_model_dir=os.path.join(temp_dir, 'saved_model'))
-
-    tflite_model = export_tflite_lib.convert_tflite_model(
-        saved_model_dir=os.path.join(temp_dir, 'saved_model'),
-        quant_type=quant_type,
-        params=params,
-        calibration_steps=5)
-
-    self.assertIsInstance(tflite_model, bytes)
-
-if __name__ == '__main__':
-  tf.test.main()