Update code to v2.11.0

official/projects/edgetpu/vision/serving/tflite_imagenet_evaluator_run.py

-# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+# 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.

official/projects/edgetpu/vision/serving/tflite_imagenet_evaluator_test.py

-# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+# 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.

official/projects/edgetpu/vision/tasks/__init__.py

-# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+# 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.

official/projects/edgetpu/vision/tasks/image_classification.py

-# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+# 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.
@@ -28,8 +28,8 @@ from official.projects.edgetpu.vision.configs import mobilenet_edgetpu_config as
 from official.projects.edgetpu.vision.dataloaders import classification_input
 from official.projects.edgetpu.vision.modeling import mobilenet_edgetpu_v1_model
 from official.projects.edgetpu.vision.modeling import mobilenet_edgetpu_v2_model
-from official.vision.beta.configs import image_classification as base_cfg
-from official.vision.beta.dataloaders import input_reader_factory
+from official.vision.configs import image_classification as base_cfg
+from official.vision.dataloaders import input_reader_factory
 
 
 def _copy_recursively(src: str, dst: str) -> None:
@@ -265,7 +265,7 @@ class EdgeTPUTask(base_task.Task):
     """Does forward and backward.
 
     Args:
-      inputs: A tuple of of input tensors of (features, labels).
+      inputs: A tuple of input tensors of (features, labels).
       model: A tf.keras.Model instance.
       optimizer: The optimizer for this training step.
       metrics: A nested structure of metrics objects.
@@ -319,7 +319,7 @@ class EdgeTPUTask(base_task.Task):
     """Runs validatation step.
 
     Args:
-      inputs: A tuple of of input tensors of (features, labels).
+      inputs: A tuple of input tensors of (features, labels).
       model: A tf.keras.Model instance.
       metrics: A nested structure of metrics objects.
 

official/projects/edgetpu/vision/tasks/image_classification_test.py

-# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+# 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.
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Lint as: python3
 """Tests for image classification task."""
 
 # pylint: disable=unused-import
@@ -20,11 +19,11 @@ from absl.testing import parameterized
 import orbit
 import tensorflow as tf
 
-from official.common import registry_imports
 from official.core import exp_factory
 from official.modeling import optimization
 from official.projects.edgetpu.vision.configs import mobilenet_edgetpu_config
 from official.projects.edgetpu.vision.tasks import image_classification
+from official.vision import registry_imports
 
 
 # Dummy ImageNet TF dataset.

official/projects/edgetpu/vision/tasks/semantic_segmentation.py

-# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+# 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.
@@ -27,11 +27,11 @@ from official.projects.edgetpu.vision.modeling import mobilenet_edgetpu_v1_model
 from official.projects.edgetpu.vision.modeling import mobilenet_edgetpu_v2_model
 from official.projects.edgetpu.vision.modeling.backbones import mobilenet_edgetpu  # pylint: disable=unused-import
 from official.projects.edgetpu.vision.modeling.heads import bifpn_head
-from official.vision.beta.dataloaders import input_reader_factory
-from official.vision.beta.dataloaders import segmentation_input
-from official.vision.beta.dataloaders import tfds_factory
-from official.vision.beta.ops import preprocess_ops
-from official.vision.beta.tasks import semantic_segmentation
+from official.vision.dataloaders import input_reader_factory
+from official.vision.dataloaders import segmentation_input
+from official.vision.dataloaders import tfds_factory
+from official.vision.ops import preprocess_ops
+from official.vision.tasks import semantic_segmentation
 
 
 class ClassMappingParser(segmentation_input.Parser):

official/projects/edgetpu/vision/tasks/semantic_segmentation_test.py

-# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+# 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.
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Lint as: python3
 """Tests for semantic segmentation task."""
 
 # pylint: disable=unused-import
@@ -20,12 +19,12 @@ from absl.testing import parameterized
 import orbit
 import tensorflow as tf
 
+from official import vision
 from official.core import exp_factory
 from official.modeling import optimization
 from official.projects.edgetpu.vision.configs import semantic_segmentation_config as seg_cfg
 from official.projects.edgetpu.vision.configs import semantic_segmentation_searched_config as autoseg_cfg
 from official.projects.edgetpu.vision.tasks import semantic_segmentation as img_seg_task
-from official.vision import beta
 
 
 # Dummy ADE20K TF dataset.

official/projects/edgetpu/vision/train.py

-# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+# 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.
@@ -12,16 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Lint as: python3
 """TensorFlow Model Garden Vision training for MobileNet-EdgeTPU."""
 
 from absl import app
 from absl import flags
 import gin
 
-# pylint: disable=unused-import
-from official.common import registry_imports
-# pylint: enable=unused-import
 from official.common import distribute_utils
 from official.common import flags as tfm_flags
 from official.core import task_factory
@@ -35,6 +31,7 @@ from official.projects.edgetpu.vision.configs import semantic_segmentation_searc
 from official.projects.edgetpu.vision.modeling.backbones import mobilenet_edgetpu
 from official.projects.edgetpu.vision.tasks import image_classification
 from official.projects.edgetpu.vision.tasks import semantic_segmentation
+from official.vision 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