Merge branch 'master' into panoptic-deeplab

official/vision/beta/projects/deepmac_maskrcnn/modeling/heads/instance_heads.py → official/projects/deepmac_maskrcnn/modeling/heads/instance_heads.py

@@ -20,7 +20,7 @@ from absl import logging
 import tensorflow as tf
 
 from official.modeling import tf_utils
-from official.vision.beta.projects.deepmac_maskrcnn.modeling.heads import hourglass_network
+from official.projects.deepmac_maskrcnn.modeling.heads import hourglass_network
 
 
 class DeepMaskHead(tf.keras.layers.Layer):

official/vision/beta/projects/deepmac_maskrcnn/modeling/heads/instance_heads_test.py → official/projects/deepmac_maskrcnn/modeling/heads/instance_heads_test.py

@@ -20,7 +20,7 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from official.vision.beta.projects.deepmac_maskrcnn.modeling.heads import instance_heads as deep_instance_heads
+from official.projects.deepmac_maskrcnn.modeling.heads import instance_heads as deep_instance_heads
 
 
 class MaskHeadTest(parameterized.TestCase, tf.test.TestCase):

official/vision/beta/projects/deepmac_maskrcnn/modeling/maskrcnn_model_test.py → official/projects/deepmac_maskrcnn/modeling/maskrcnn_model_test.py

@@ -21,8 +21,8 @@ from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf
 
-from official.vision.beta.projects.deepmac_maskrcnn.modeling import maskrcnn_model
-from official.vision.beta.projects.deepmac_maskrcnn.modeling.heads import instance_heads as deep_instance_heads
+from official.projects.deepmac_maskrcnn.modeling import maskrcnn_model
+from official.projects.deepmac_maskrcnn.modeling.heads import instance_heads as deep_instance_heads
 from official.vision.modeling.backbones import resnet
 from official.vision.modeling.decoders import fpn
 from official.vision.modeling.heads import dense_prediction_heads

official/vision/beta/projects/deepmac_maskrcnn/serving/detection.py → official/projects/deepmac_maskrcnn/serving/detection.py

@@ -19,9 +19,9 @@ from typing import Dict, Mapping, Text
 
 import tensorflow as tf
 
-from official.vision.beta.projects.deepmac_maskrcnn.configs import deep_mask_head_rcnn as cfg
-from official.vision.beta.projects.deepmac_maskrcnn.modeling import maskrcnn_model
-from official.vision.beta.projects.deepmac_maskrcnn.tasks import deep_mask_head_rcnn
+from official.projects.deepmac_maskrcnn.configs import deep_mask_head_rcnn as cfg
+from official.projects.deepmac_maskrcnn.modeling import maskrcnn_model
+from official.projects.deepmac_maskrcnn.tasks import deep_mask_head_rcnn
 from official.vision.ops import box_ops
 from official.vision.serving import detection
 

official/vision/beta/projects/deepmac_maskrcnn/serving/detection_test.py → official/projects/deepmac_maskrcnn/serving/detection_test.py

@@ -24,7 +24,7 @@ from PIL import Image
 import tensorflow as tf
 
 from official.core import exp_factory
-from official.vision.beta.projects.deepmac_maskrcnn.serving import detection
+from official.projects.deepmac_maskrcnn.serving import detection
 
 
 class DetectionExportTest(tf.test.TestCase, parameterized.TestCase):

official/vision/beta/projects/deepmac_maskrcnn/serving/export_saved_model.py → official/projects/deepmac_maskrcnn/serving/export_saved_model.py

@@ -38,8 +38,8 @@ from absl import flags
 
 from official.core import exp_factory
 from official.modeling import hyperparams
-from official.vision.beta.projects.deepmac_maskrcnn.serving import detection
-from official.vision.beta.projects.deepmac_maskrcnn.tasks import deep_mask_head_rcnn  # pylint: disable=unused-import
+from official.projects.deepmac_maskrcnn.serving import detection
+from official.projects.deepmac_maskrcnn.tasks import deep_mask_head_rcnn  # pylint: disable=unused-import
 from official.vision.serving import export_saved_model_lib
 
 FLAGS = flags.FLAGS

official/vision/beta/projects/deepmac_maskrcnn/tasks/deep_mask_head_rcnn.py → official/projects/deepmac_maskrcnn/tasks/deep_mask_head_rcnn.py

@@ -17,9 +17,9 @@
 import tensorflow as tf
 
 from official.core import task_factory
-from official.vision.beta.projects.deepmac_maskrcnn.configs import deep_mask_head_rcnn as deep_mask_head_rcnn_config
-from official.vision.beta.projects.deepmac_maskrcnn.modeling import maskrcnn_model as deep_maskrcnn_model
-from official.vision.beta.projects.deepmac_maskrcnn.modeling.heads import instance_heads as deep_instance_heads
+from official.projects.deepmac_maskrcnn.configs import deep_mask_head_rcnn as deep_mask_head_rcnn_config
+from official.projects.deepmac_maskrcnn.modeling import maskrcnn_model as deep_maskrcnn_model
+from official.projects.deepmac_maskrcnn.modeling.heads import instance_heads as deep_instance_heads
 from official.vision.modeling import backbones
 from official.vision.modeling.decoders import factory as decoder_factory
 from official.vision.modeling.heads import dense_prediction_heads

official/vision/beta/projects/deepmac_maskrcnn/train.py → official/projects/deepmac_maskrcnn/train.py

@@ -28,7 +28,7 @@ from official.core import train_lib
 from official.core import train_utils
 from official.modeling import performance
 # pylint: disable=unused-import
-from official.vision.beta.projects.deepmac_maskrcnn.common import registry_imports
+from official.projects.deepmac_maskrcnn.common import registry_imports
 # pylint: enable=unused-import
 
 FLAGS = flags.FLAGS

official/projects/labse/README.md

+# Language-agnostic BERT Sentence Embedding
+
+The repository contains the implementation and experiment definition of `LaBSE`,
+[Language-agnostic BERT Sentence Embedding](https://arxiv.org/pdf/2007.01852.pdf).
+The implementation is provided by the paper author, Yinfei Yang. Note that,
+the cross-accelerator batch softmax is not implemented by the author, so the
+implementation does not fully reproduce the paper yet.
+
+Due to the data policy, the authors are not able to release the pre-training and
+fine-tuning data for `LaBSE` training.
+
+### Requirements
+
+The starter code requires Tensorflow. If you haven't installed it yet, follow
+the instructions on [tensorflow.org][1].
+This code has been tested with Tensorflow 2.8.0. Going forward,
+we will continue to target the latest released version of Tensorflow.
+
+Please verify that you have Python 3.7+ and Tensorflow 2.8.0 or higher
+installed by running the following commands:
+
+```sh
+python --version
+python -c 'import tensorflow as tf; print(tf.__version__)'
+```
+
+Refer to the [instructions here][2]
+for using the model in this repo. Make sure to add the models folder to your
+Python path.
+
+[1]: https://www.tensorflow.org/install/
+[2]:
+https://github.com/tensorflow/models/tree/master/official#running-the-models
+
+## Data
+
+The pre-training data should be multi-lingual and the format is the same as BERT
+pre-training.
+
+The fine-tuning data follows the format as below:
+
+```text
+{   # (tensorflow.Example)
+  features: {
+    feature: {
+      key  : "src_raw"
+      value: {
+        bytes_list: {
+          value: [ "Foo. " ]
+        }
+      }
+    }
+    feature: {
+      key  : "tgt_raw"
+      value: {
+        bytes_list: {
+          value: [ "Bar. " ]
+        }
+      }
+    }
+  }
+}
+```
+
+## Train using the config file.
+
+After you generated your pretraining data, run the following command to start
+pretraining:
+
+```bash
+TPU=local
+VOCAB=???
+INIT_CHECKPOINT=???
+PARAMS="task.train_data.input_data=/path/to/train/data"
+PARAMS="${PARAMS},task.train_data.vocab_file=${VOCAB}"
+PARAMS="${PARAMS},task.validation_data.input_path=/path/to/validation/data"
+PARAMS="${PARAMS},task.validation_data.vocab_file=${VOCAB}"
+PARAMS="${PARAMS},task.init_checkpoint=${INIT_CHECKPOINT}"
+PARAMS="${PARAMS},runtime.distribution_strategy=tpu"
+
+python3 train.py \
+  --experiment=labse/train \
+  --config_file=./experiments/labse_bert_base.yaml \
+  --config_file=./experiments/labse_base.yaml \
+  --params_override=${PARAMS} \
+  --tpu=${TPU} \
+  --model_dir=/folder/to/hold/logs/and/models/ \
+  --mode=train_and_eval
+```
+
+## Implementation
+
+We implement the encoder and layers using `tf.keras` APIs in NLP
+modeling library:
+
+  * [dual_encoder.py](https://github.com/tensorflow/models/blob/master/official/nlp/tasks/dual_encoder.py)
+  contains the dual-encoder task used for labse training.
+
+  * [config_labse.py](https://github.com/tensorflow/models/blob/master/official/projects/labse/config_labse.py)
+  registers the labse training experiment.
+
+  * [train.py](https://github.com/tensorflow/models/blob/master/official/projects/labse/train.py)
+  is the program entry.
+
+
+## Pre-trained model through TF-HUB
+
+If you are looking for pre-trained models, please check out:
+https://tfhub.dev/google/LaBSE/2.
+The hub `SavedModel`s are exported through the `export_tfhub.py` in
+this repository.

official/projects/labse/config_labse.py

+# Copyright 2022 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.
+
+# pylint: disable=g-doc-return-or-yield,line-too-long
+"""LaBSE configurations."""
+import dataclasses
+from official.core import config_definitions as cfg
+from official.core import exp_factory
+from official.modeling import optimization
+from official.nlp.data import dual_encoder_dataloader
+from official.nlp.tasks import dual_encoder
+
+AdamWeightDecay = optimization.AdamWeightDecayConfig
+PolynomialLr = optimization.PolynomialLrConfig
+PolynomialWarmupConfig = optimization.PolynomialWarmupConfig
+
+
+@dataclasses.dataclass
+class LaBSEOptimizationConfig(optimization.OptimizationConfig):
+  """Bert optimization config."""
+  optimizer: optimization.OptimizerConfig = optimization.OptimizerConfig(
+      type="adamw", adamw=AdamWeightDecay())
+  learning_rate: optimization.LrConfig = optimization.LrConfig(
+      type="polynomial",
+      polynomial=PolynomialLr(
+          initial_learning_rate=1e-4,
+          decay_steps=1000000,
+          end_learning_rate=0.0))
+  warmup: optimization.WarmupConfig = optimization.WarmupConfig(
+      type="polynomial", polynomial=PolynomialWarmupConfig(warmup_steps=10000))
+
+
+@exp_factory.register_config_factory("labse/train")
+def labse_train() -> cfg.ExperimentConfig:
+  r"""Language-agnostic bert sentence embedding.
+
+  *Note*: this experiment does not use cross-accelerator global softmax so it
+  does not reproduce the exact LABSE training.
+  """
+  config = cfg.ExperimentConfig(
+      task=dual_encoder.DualEncoderConfig(
+          train_data=dual_encoder_dataloader.DualEncoderDataConfig(),
+          validation_data=dual_encoder_dataloader.DualEncoderDataConfig(
+              is_training=False, drop_remainder=False)),
+      trainer=cfg.TrainerConfig(
+          optimizer_config=LaBSEOptimizationConfig(
+              learning_rate=optimization.LrConfig(
+                  type="polynomial",
+                  polynomial=PolynomialLr(
+                      initial_learning_rate=3e-5, end_learning_rate=0.0)),
+              warmup=optimization.WarmupConfig(
+                  type="polynomial", polynomial=PolynomialWarmupConfig()))),
+      restrictions=[
+          "task.train_data.is_training != None",
+          "task.validation_data.is_training != None"
+      ])
+  return config

official/projects/labse/experiments/labse_base.yaml

+task:
+  hub_module_url: ''
+  model:
+    bidirectional: true
+    max_sequence_length: 32
+    logit_scale: 100
+    logit_margin: 0.3
+  init_checkpoint: 'the pre-trained BERT checkpoint using the labse vocab.'
+  train_data:
+    drop_remainder: true
+    global_batch_size: 4096
+    input_path: 'the path to train partition'
+    left_text_fields: ['src_raw']
+    right_text_fields: ['tgt_raw']
+    vocab_file: 'the path to vocab.txt'
+    lower_case: false
+    is_training: true
+    seq_length: 32
+    sharding: false
+    cycle_length: 4
+    shuffle_buffer_size: 1000
+    tfds_as_supervised: false
+    tfds_data_dir: ''
+    tfds_name: ''
+    tfds_skip_decoding_feature: ''
+    tfds_split: ''
+  validation_data:
+    block_length: 1
+    cache: false
+    cycle_length: 4
+    drop_remainder: false
+    global_batch_size: 32000
+    input_path: 'the path to validation partition'
+    left_text_fields: ['src_raw']
+    right_text_fields: ['tgt_raw']
+    vocab_file: 'the path to vocab.txt'
+    lower_case: false
+    is_training: false
+    seq_length: 32
+    sharding: true
+    shuffle_buffer_size: 1000
+    tfds_as_supervised: false
+    tfds_data_dir: ''
+    tfds_name: ''
+    tfds_skip_decoding_feature: ''
+    tfds_split: ''
+trainer:
+  checkpoint_interval: 1000
+  eval_tf_function: true
+  max_to_keep: 5
+  optimizer_config:
+    learning_rate:
+      polynomial:
+        cycle: false
+        decay_steps: 500000
+        end_learning_rate: 0.0
+        initial_learning_rate: 1.0e-04
+        name: PolynomialDecay
+        power: 1.0
+      type: polynomial
+    optimizer:
+      adamw:
+        amsgrad: false
+        beta_1: 0.9
+        beta_2: 0.999
+        epsilon: 1.0e-05
+        exclude_from_weight_decay: null
+        include_in_weight_decay: null
+        name: AdamWeightDecay
+        weight_decay_rate: 0.0
+        gradient_clip_norm: 100
+      type: adamw
+    warmup:
+      polynomial:
+        name: polynomial
+        power: 1
+        warmup_steps: 5000
+      type: polynomial
+  steps_per_loop: 1000
+  summary_interval: 1000
+  train_tf_function: true
+  train_tf_while_loop: true
+  train_steps: 500000
+  validation_interval: 1000
+  validation_steps: 100

official/projects/labse/experiments/labse_bert_base.yaml

+task:
+  model:
+    encoder:
+      bert:
+        attention_dropout_rate: 0.1
+        dropout_rate: 0.1
+        hidden_activation: gelu
+        hidden_size: 768
+        initializer_range: 0.02
+        intermediate_size: 3072
+        max_position_embeddings: 512
+        num_attention_heads: 12
+        num_layers: 12
+        type_vocab_size: 2
+        vocab_size: 501153

official/projects/labse/export_tfhub.py

+# Copyright 2022 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.
+
+r"""Exports the LaBSE model and its preprocessing as SavedModels for TF Hub.
+
+Example usage:
+# Point this variable to your training results.
+# Note that flag --do_lower_case is inferred from the name.
+LaBSE_DIR=<Your LaBSE model dir>
+# Step 1: export the core LaBSE model.
+python3 ./export_tfhub.py \
+  --bert_config_file ${LaBSE_DIR:?}/bert_config.json \
+  --model_checkpoint_path ${LaBSE_DIR:?}/labse_model.ckpt \
+  --vocab_file ${LaBSE_DIR:?}/vocab.txt \
+  --export_type model --export_path /tmp/labse_model
+# Step 2: export matching preprocessing (be sure to use same flags).
+python3 ./export_tfhub.py \
+  --vocab_file ${LaBSE_DIR:?}/vocab.txt \
+  --export_type preprocessing --export_path /tmp/labse_preprocessing
+"""
+
+from typing import Text
+
+from absl import app
+from absl import flags
+from absl import logging
+import tensorflow as tf
+
+from official.legacy.bert import bert_models
+from official.legacy.bert import configs
+from official.nlp.modeling import models
+from official.nlp.tasks import utils
+from official.nlp.tools import export_tfhub_lib
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_enum("export_type", "model", ["model", "preprocessing"],
+                  "The type of model to export")
+flags.DEFINE_string("export_path", None, "TF-Hub SavedModel destination path.")
+flags.DEFINE_string(
+    "bert_tfhub_module", None,
+    "Bert tfhub module to define core bert layers. Needed for --export_type "
+    "model.")
+flags.DEFINE_string(
+    "bert_config_file", None,
+    "Bert configuration file to define core bert layers. It will not be used "
+    "if bert_tfhub_module is set. Needed for --export_type model.")
+flags.DEFINE_string(
+    "model_checkpoint_path", None, "File path to TF model checkpoint. "
+    "Needed for --export_type model.")
+flags.DEFINE_string(
+    "vocab_file", None,
+    "The vocabulary file that the BERT model was trained on. "
+    "Needed for both --export_type model and preprocessing.")
+flags.DEFINE_bool(
+    "do_lower_case", None,
+    "Whether to lowercase before tokenization. If left as None, "
+    "do_lower_case will be enabled if 'uncased' appears in the "
+    "name of --vocab_file. "
+    "Needed for both --export_type model and preprocessing.")
+flags.DEFINE_integer(
+    "default_seq_length", 128,
+    "The sequence length of preprocessing results from "
+    "top-level preprocess method. This is also the default "
+    "sequence length for the bert_pack_inputs subobject."
+    "Needed for --export_type preprocessing.")
+flags.DEFINE_bool(
+    "tokenize_with_offsets", False,  # TODO(b/181866850)
+    "Whether to export a .tokenize_with_offsets subobject for "
+    "--export_type preprocessing.")
+flags.DEFINE_bool(
+    "normalize", True,
+    "Parameter of DualEncoder model, normalize the embedding (pooled_output) "
+    "if set to True.")
+
+
+def _get_do_lower_case(do_lower_case, vocab_file):
+  """Returns do_lower_case, replacing None by a guess from vocab file name."""
+  if do_lower_case is None:
+    do_lower_case = "uncased" in vocab_file
+    logging.info("Using do_lower_case=%s based on name of vocab_file=%s",
+                 do_lower_case, vocab_file)
+  return do_lower_case
+
+
+def create_labse_model(bert_tfhub_module: Text,
+                       bert_config: configs.BertConfig,
+                       normalize: bool) -> tf.keras.Model:
+  """Creates a LaBSE keras core model from BERT configuration.
+
+  Args:
+    bert_tfhub_module: The bert tfhub module path. The LaBSE will be built upon
+      the tfhub module if it is not empty.
+    bert_config: A `BertConfig` to create the core model. Used if
+      bert_tfhub_module is empty.
+    normalize: Parameter of DualEncoder model, normalize the embedding (
+      pooled_output) if set to True.
+
+  Returns:
+    A keras model.
+  """
+  if bert_tfhub_module:
+    encoder_network = utils.get_encoder_from_hub(bert_tfhub_module)
+  else:
+    encoder_network = bert_models.get_transformer_encoder(
+        bert_config, sequence_length=None)
+
+  labse_model = models.DualEncoder(
+      network=encoder_network,
+      max_seq_length=None,
+      normalize=normalize,
+      output="predictions")
+  return labse_model, encoder_network  # pytype: disable=bad-return-type  # typed-keras
+
+
+def export_labse_model(bert_tfhub_module: Text, bert_config: configs.BertConfig,
+                       model_checkpoint_path: Text, hub_destination: Text,
+                       vocab_file: Text, do_lower_case: bool, normalize: bool):
+  """Restores a tf.keras.Model and saves for TF-Hub."""
+  core_model, encoder = create_labse_model(
+      bert_tfhub_module, bert_config, normalize)
+  checkpoint = tf.train.Checkpoint(encoder=encoder)
+  checkpoint.restore(model_checkpoint_path).assert_existing_objects_matched()
+  core_model.vocab_file = tf.saved_model.Asset(vocab_file)
+  core_model.do_lower_case = tf.Variable(do_lower_case, trainable=False)
+  core_model.save(hub_destination, include_optimizer=False, save_format="tf")
+
+
+def main(_):
+  do_lower_case = export_tfhub_lib.get_do_lower_case(FLAGS.do_lower_case,
+                                                     FLAGS.vocab_file)
+  if FLAGS.export_type == "model":
+    if FLAGS.bert_tfhub_module:
+      bert_config = None
+    else:
+      bert_config = configs.BertConfig.from_json_file(FLAGS.bert_config_file)
+    export_labse_model(FLAGS.bert_tfhub_module, bert_config,
+                       FLAGS.model_checkpoint_path, FLAGS.export_path,
+                       FLAGS.vocab_file, do_lower_case, FLAGS.normalize)
+  elif FLAGS.export_type == "preprocessing":
+    # LaBSE is still a BERT model, reuse the export_bert_preprocessing here.
+    export_tfhub_lib.export_bert_preprocessing(
+        FLAGS.export_path, FLAGS.vocab_file, do_lower_case,
+        FLAGS.default_seq_length, FLAGS.tokenize_with_offsets)
+  else:
+    raise app.UsageError("Unknown value '%s' for flag --export_type")
+
+
+if __name__ == "__main__":
+  app.run(main)

official/projects/labse/export_tfhub_test.py

+# Copyright 2022 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 LaBSE's export_tfhub."""
+
+import os
+
+# Import libraries
+import numpy as np
+import tensorflow as tf
+import tensorflow_hub as hub
+from official.legacy.bert import configs
+from official.projects.labse import export_tfhub
+
+
+class ExportModelTest(tf.test.TestCase):
+
+  def test_export_model(self):
+    # Exports a savedmodel for TF-Hub
+    hidden_size = 16
+    bert_config = configs.BertConfig(
+        vocab_size=100,
+        hidden_size=hidden_size,
+        intermediate_size=32,
+        max_position_embeddings=128,
+        num_attention_heads=2,
+        num_hidden_layers=1)
+    labse_model, encoder = export_tfhub.create_labse_model(
+        None, bert_config, normalize=True)
+    model_checkpoint_dir = os.path.join(self.get_temp_dir(), "checkpoint")
+    checkpoint = tf.train.Checkpoint(encoder=encoder)
+    checkpoint.save(os.path.join(model_checkpoint_dir, "test"))
+    model_checkpoint_path = tf.train.latest_checkpoint(model_checkpoint_dir)
+
+    vocab_file = os.path.join(self.get_temp_dir(), "uncased_vocab.txt")
+    with tf.io.gfile.GFile(vocab_file, "w") as f:
+      f.write("dummy content")
+
+    hub_destination = os.path.join(self.get_temp_dir(), "hub")
+    export_tfhub.export_labse_model(
+        None,  # bert_tfhub_module
+        bert_config,
+        model_checkpoint_path,
+        hub_destination,
+        vocab_file,
+        do_lower_case=True,
+        normalize=True)
+
+    # Restores a hub KerasLayer.
+    hub_layer = hub.KerasLayer(hub_destination, trainable=True)
+
+    if hasattr(hub_layer, "resolved_object"):
+      # Checks meta attributes.
+      self.assertTrue(hub_layer.resolved_object.do_lower_case.numpy())
+      with tf.io.gfile.GFile(
+          hub_layer.resolved_object.vocab_file.asset_path.numpy()) as f:
+        self.assertEqual("dummy content", f.read())
+    # Checks the hub KerasLayer.
+    for source_weight, hub_weight in zip(labse_model.trainable_weights,
+                                         hub_layer.trainable_weights):
+      self.assertAllClose(source_weight.numpy(), hub_weight.numpy())
+
+    seq_length = 10
+    dummy_ids = np.zeros((2, seq_length), dtype=np.int32)
+    hub_outputs = hub_layer([dummy_ids, dummy_ids, dummy_ids])
+    source_outputs = labse_model([dummy_ids, dummy_ids, dummy_ids])
+
+    self.assertEqual(hub_outputs["pooled_output"].shape, (2, hidden_size))
+    self.assertEqual(hub_outputs["sequence_output"].shape,
+                     (2, seq_length, hidden_size))
+    for output_name in source_outputs:
+      self.assertAllClose(hub_outputs[output_name].numpy(),
+                          hub_outputs[output_name].numpy())
+
+    # Test that training=True makes a difference (activates dropout).
+    def _dropout_mean_stddev(training, num_runs=20):
+      input_ids = np.array([[14, 12, 42, 95, 99]], np.int32)
+      inputs = [input_ids, np.ones_like(input_ids), np.zeros_like(input_ids)]
+      outputs = np.concatenate([
+          hub_layer(inputs, training=training)["pooled_output"]
+          for _ in range(num_runs)
+      ])
+      return np.mean(np.std(outputs, axis=0))
+
+    self.assertLess(_dropout_mean_stddev(training=False), 1e-6)
+    self.assertGreater(_dropout_mean_stddev(training=True), 1e-3)
+
+    # Test propagation of seq_length in shape inference.
+    input_word_ids = tf.keras.layers.Input(shape=(seq_length,), dtype=tf.int32)
+    input_mask = tf.keras.layers.Input(shape=(seq_length,), dtype=tf.int32)
+    input_type_ids = tf.keras.layers.Input(shape=(seq_length,), dtype=tf.int32)
+    outputs = hub_layer([input_word_ids, input_mask, input_type_ids])
+    self.assertEqual(outputs["pooled_output"].shape.as_list(),
+                     [None, hidden_size])
+    self.assertEqual(outputs["sequence_output"].shape.as_list(),
+                     [None, seq_length, hidden_size])
+
+
+if __name__ == "__main__":
+  tf.test.main()

official/projects/labse/train.py

+# Copyright 2022 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.
+
+"""TensorFlow Model Garden Labse training driver, register labse configs."""
+
+# pylint: disable=unused-import
+from absl import app
+
+from official.common import flags as tfm_flags
+from official.nlp import tasks
+from official.nlp import train
+from official.projects.labse import config_labse
+
+if __name__ == '__main__':
+  tfm_flags.define_flags()
+  app.run(train.main)

official/projects/movinet/README.md

@@ -8,12 +8,17 @@ This repository is the official implementation of
 [MoViNets: Mobile Video Networks for Efficient Video
 Recognition](https://arxiv.org/abs/2103.11511).
 
-**[UPDATE 2021-07-12] Mobile Models Available via [TF Lite](#tf-lite-streaming-models)**
+-   **[UPDATE 2022-03-14] Quantized TF Lite models
+    [available on TF Hub](https://tfhub.dev/s?deployment-format=lite&q=movinet)
+    (also [see table](https://tfhub.dev/google/collections/movinet) for
+    quantized performance)**
 
 <p align="center">
   <img src="https://storage.googleapis.com/tf_model_garden/vision/movinet/artifacts/hoverboard_stream.gif" height=500>
 </p>
 
+Create your own video plot like the one above with this [Colab notebook](https://colab.research.google.com/github/tensorflow/models/blob/master/official/projects/movinet/tools/plot_movinet_video_stream_predictions.ipynb).
+
 ## Description
 
 Mobile Video Networks (MoViNets) are efficient video classification models
@@ -55,6 +60,8 @@ approach that performs redundant computation and limits temporal scope.
 
 ## History
 
+- **2022-03-14** Support quantized TF Lite models and add/update Colab
+notebooks.
 - **2021-07-12** Add TF Lite support and replace 3D stream models with
 mobile-friendly (2+1)D stream.
 - **2021-05-30** Add streaming MoViNet checkpoints and examples.
@@ -71,6 +78,7 @@ mobile-friendly (2+1)D stream.
 - [Requirements](#requirements)
 - [Results and Pretrained Weights](#results-and-pretrained-weights)
   - [Kinetics 600](#kinetics-600)
+  - [Kinetics 400](#kinetics-400)
 - [Prediction Examples](#prediction-examples)
 - [TF Lite Example](#tf-lite-example)
 - [Training and Evaluation](#training-and-evaluation)
@@ -165,7 +173,7 @@ different architecture. To download the old checkpoints, insert `_legacy` before
 
 For convenience, we provide converted TF Lite models for inference on mobile
 devices. See the [TF Lite Example](#tf-lite-example) to export and run your own
-models.
+models. We also provide [quantized TF Lite binaries via TF Hub](https://tfhub.dev/s?deployment-format=lite&q=movinet).
 
 For reference, MoViNet-A0-Stream runs with a similar latency to
 [MobileNetV3-Large]
@@ -188,7 +196,7 @@ W-2135 CPU.
 
 ### Kinetics 400
 
-We also have checkpoints for Kinetics 600 models available. See the Kinetics 600
+We also have checkpoints for Kinetics 400 models available. See the Kinetics 600
 sections for more details. To load checkpoints, set `num_classes=400`.
 
 #### Base Models
@@ -226,7 +234,7 @@ backbone = movinet.Movinet(
     use_external_states=False,
 )
 model = movinet_model.MovinetClassifier(
-    backbone, num_classes=600, output_states=True)
+    backbone, num_classes=600, output_states=False)
 
 # Create your example input here.
 # Refer to the paper for recommended input shapes.