Internal change

PiperOrigin-RevId: 428078415

official/projects/qat/vision/README.md

+# Quantization Aware Training Project for Computer Vision Models
+
+[TOC]
+
+⚠️ Disclaimer: All datasets hyperlinked from this page are not owned or
+distributed by Google. The dataset is made available by third parties.
+Please review the terms and conditions made available by the third parties
+before using the data.
+
+## Overview
+
+This project includes quantization aware training code for Computer Vision
+models. These are examples to show how to apply the Model Optimization Toolkit's
+[quantization aware training API](https://www.tensorflow.org/model_optimization/guide/quantization/training).
+
+Note: Currently, we support a limited number of ML tasks & models (e.g., image
+classification and semantic segmentation)
+We will keep adding support for other ML tasks and models in the next releases.
+
+## How to train a model
+
+```
+EXPERIMENT=xxx  # Change this for your run, for example, 'mobilenet_imagenet_qat'
+CONFIG_FILE=xxx  # Change this for your run, for example, path of imagenet_mobilenetv2_qat_gpu.yaml
+MODEL_DIR=xxx  #  Change this for your run, for example, /tmp/model_dir
+$ python3 train.py \
+--experiment=${EXPERIMENT} \
+--config_file=${CONFIG_FILE} \
+--model_dir=${MODEL_DIR} \
+--mode=train_and_eval
+```
+
+## Model Accuracy
+
+<figure align="center">
+<img width=70% src=https://storage.googleapis.com/tf_model_garden/models/qat/images/readme-qat-classification-plot.png>
+<figcaption>Comparison of Imagenet top-1 accuracy for the classification models</figcaption>
+</figure>
+
+Note: The Top-1 model accuracy is measured on the validation set of [ImageNet](https://www.image-net.org/).
+
+
+### Pre-trained Models
+
+|Model                 |Resolution|Top-1 Accuracy (FP32)|Top-1 Accuracy (Int8/PTQ)|Top-1 Accuracy (Int8/QAT)|Config                                                                                                                                                              |Download                                                                                                                                        |
+|----------------------|----------|---------------------|-------------------------|-------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------|
+|MobileNetV2           |224x224   |72.782%              |72.392%                  |72.792%                  |[config](https://github.com/tensorflow/models/blob/master/official/projects/qat/vision/configs/experiments/image_classification/imagenet_mobilenetv2_qat_gpu.yaml)  |[TFLite(Int8/QAT)](https://storage.googleapis.com/tf_model_garden/vision/mobilenet/v2_1.0_int8/mobilenet_v2_1.00_224_int8.tflite)                    |
+|ResNet50              |224x224   |76.710%              |76.420%                  |77.200%                  |[config](https://github.com/tensorflow/models/blob/master/official/projects/qat/vision/configs/experiments/image_classification/imagenet_resnet50_qat_gpu.yaml)     |[TFLite(Int8/QAT)](https://storage.googleapis.com/tf_model_garden/vision/resnet50_imagenet/resnet_50_224_int8.tflite)                                |
+|MobileNetV3.5 MultiAVG|224x224   |75.212%              |74.122%                  |75.130%                  |[config](https://github.com/tensorflow/models/blob/master/official/projects/qat/vision/configs/experiments/image_classification/imagenet_mobilenetv3.5_qat_gpu.yaml)|[TFLite(Int8/QAT)](https://storage.googleapis.com/tf_model_garden/vision/mobilenet/v3.5multiavg_1.0_int8/mobilenet_v3.5multiavg_1.00_224_int8.tflite)|
+

official/projects/qat/vision/configs/__init__.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.
+
+# Lint as: python3
+"""Configs package definition."""
+
+from official.projects.qat.vision.configs import image_classification
+from official.projects.qat.vision.configs import semantic_segmentation

official/projects/qat/vision/configs/common.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.
+
+# Lint as: python3
+"""Image classification configuration definition."""
+
+import dataclasses
+from typing import Optional
+
+from official.modeling import hyperparams
+
+
+@dataclasses.dataclass
+class Quantization(hyperparams.Config):
+  """Quantization parameters.
+
+  Attributes:
+    pretrained_original_checkpoint: A string indicate pretrained checkpoint
+      location.
+    change_num_bits: A `bool` indicates whether to manually allocate num_bits.
+    num_bits_weight: An `int` number of bits for weight. Default to 8.
+    num_bits_activation: An `int` number of bits for activation. Default to 8.
+  """
+  pretrained_original_checkpoint: Optional[str] = None
+  change_num_bits: bool = False
+  num_bits_weight: int = 8
+  num_bits_activation: int = 8

official/projects/qat/vision/configs/experiments/image_classification/imagenet_mobilenetv2_qat_gpu.yaml

+runtime:
+  distribution_strategy: 'mirrored'
+  mixed_precision_dtype: 'float32'
+  loss_scale: 'dynamic'
+task:
+  model:
+    num_classes: 1001
+    input_size: [224, 224, 3]
+    backbone:
+      type: 'mobilenet'
+      mobilenet:
+        model_id: 'MobileNetV2'
+        filter_size_scale: 1.0
+    dropout_rate: 0.1
+  losses:
+    l2_weight_decay: 0.0000001
+    one_hot: true
+    label_smoothing: 0.1
+  train_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/train*'
+    is_training: true
+    global_batch_size: 512  # 64 * 8
+    dtype: 'float32'
+  validation_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/valid*'
+    is_training: false
+    global_batch_size: 512  # 64 * 8
+    dtype: 'float32'
+    drop_remainder: false
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/mobilenetv2_gpu/22984194/ckpt-625500'
+trainer:
+  # With below setting, the accuracy of QAT reaches to accuracy 0.7279 after 43 hours with 8 GPUS.
+  train_steps: 250200
+  validation_steps: 98
+  validation_interval: 2502
+  steps_per_loop: 2502
+  summary_interval: 2502
+  checkpoint_interval: 2502
+  optimizer_config:
+    learning_rate:
+      type: 'exponential'
+      exponential:
+        decay_rate: 0.9
+        decay_steps: 1251
+        initial_learning_rate: 0.0001
+        name: 'ExponentialDecay'
+        offset: 0
+        staircase: true
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 0

official/projects/qat/vision/configs/experiments/image_classification/imagenet_mobilenetv2_qat_gpu_batch256.yaml

+runtime:
+  distribution_strategy: 'mirrored'
+  mixed_precision_dtype: 'float32'
+  loss_scale: 'dynamic'
+task:
+  model:
+    num_classes: 1001
+    input_size: [224, 224, 3]
+    backbone:
+      type: 'mobilenet'
+      mobilenet:
+        model_id: 'MobileNetV2'
+        filter_size_scale: 1.0
+    dropout_rate: 0.0  # changed from 0.2 to 0.0
+  losses:
+    l2_weight_decay: 0.0000001
+    one_hot: true
+    label_smoothing: 0.1
+  train_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/train*'
+    is_training: true
+    global_batch_size: 256
+    dtype: 'float32'
+  validation_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/valid*'
+    is_training: false
+    global_batch_size: 256
+    dtype: 'float32'
+    drop_remainder: false
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/mobilenetv2_gpu/22984194/ckpt-625500'
+trainer:
+  # With below setting, the accuracy of QAT reaches Top1-accuracy 0.7251 at 420336 steps after
+  # 1 day 19 hours of training with 8GPUs, which is higher than the result of PTQ in MobileNetV2.
+  train_steps: 1000800  # 200 epochs
+  validation_steps: 196  # NUM_EXAMPLES (50000) // global_batch_size (256)
+  validation_interval: 5004  # 1 epoch
+  steps_per_loop: 5004  # NUM_EXAMPLES (1281167) // global_batch_size (256)
+  summary_interval: 5004  # 1 epoch
+  checkpoint_interval: 5004  # 1 epoch
+  max_to_keep: 200
+  optimizer_config:
+    learning_rate:
+      type: 'exponential'
+      exponential:
+        initial_learning_rate: 0.0001
+        decay_steps: 1251  # steps_per_epoch // 4
+        decay_rate: 0.96
+        staircase: true
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 0

official/projects/qat/vision/configs/experiments/image_classification/imagenet_mobilenetv2_qat_gpu_batch512.yaml

+runtime:
+  distribution_strategy: 'mirrored'
+  mixed_precision_dtype: 'float32'
+  loss_scale: 'dynamic'
+task:
+  model:
+    num_classes: 1001
+    input_size: [224, 224, 3]
+    backbone:
+      type: 'mobilenet'
+      mobilenet:
+        model_id: 'MobileNetV2'
+        filter_size_scale: 1.0
+    dropout_rate: 0.0  # changed from 0.2 to 0.0
+  losses:
+    l2_weight_decay: 0.0000001
+    one_hot: true
+    label_smoothing: 0.1
+  train_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/train*'
+    is_training: true
+    global_batch_size: 512
+    dtype: 'float32'
+  validation_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/valid*'
+    is_training: false
+    global_batch_size: 512
+    dtype: 'float32'
+    drop_remainder: false
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/mobilenetv2_gpu/22984194/ckpt-625500'
+trainer:
+  # With below setting, the accuracy of QAT reaches Top1-accuracy 0.7266 at 312750 steps after
+  # 1 day 22 hours of training with 8GPUs, which is higher than the result of PTQ in MobileNetV2.
+  train_steps: 500400  # 200 epochs
+  validation_steps: 98  # NUM_EXAMPLES (50000) // global_batch_size (512)
+  validation_interval: 2502  # 1 epoch
+  steps_per_loop: 2502  # NUM_EXAMPLES (1281167) // global_batch_size (512)
+  summary_interval: 2502  # 1 epoch
+  checkpoint_interval: 2502  # 1 epoch
+  max_to_keep: 200
+  optimizer_config:
+    learning_rate:
+      type: 'exponential'
+      exponential:
+        initial_learning_rate: 0.0002
+        decay_steps: 1251  # steps_per_epoch // 2
+        decay_rate: 0.96
+        staircase: true
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 0

official/projects/qat/vision/configs/experiments/image_classification/imagenet_mobilenetv3.5_qat_gpu.yaml

+runtime:
+  distribution_strategy: 'mirrored'
+  mixed_precision_dtype: 'float32'
+  loss_scale: 'dynamic'
+task:
+  model:
+    num_classes: 1001
+    input_size: [224, 224, 3]
+    backbone:
+      type: 'mobilenet'
+      mobilenet:
+        model_id: 'MobileNetMultiAVG'
+        filter_size_scale: 1.0
+    dropout_rate: 0.3
+  losses:
+    l2_weight_decay: 0.000001
+    one_hot: true
+    label_smoothing: 0.1
+  train_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/train*'
+    is_training: true
+    global_batch_size: 512
+    dtype: 'float32'
+  validation_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/valid*'
+    is_training: false
+    global_batch_size: 512
+    dtype: 'float32'
+    drop_remainder: false
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/tf2_mhave_nobias_bn_aug05/28334857/ckpt-156000'
+trainer:
+  # With below setting, the accuracy of QAT reaches to accuracy 0.7513 after 30 hours with 8 GPUS.
+  train_steps: 250200
+  validation_steps: 98
+  validation_interval: 2502
+  steps_per_loop: 2502
+  summary_interval: 2502
+  checkpoint_interval: 2502
+  optimizer_config:
+    learning_rate:
+      type: 'exponential'
+      exponential:
+        decay_rate: 0.9
+        decay_steps: 1251
+        initial_learning_rate: 0.0004
+        name: 'ExponentialDecay'
+        offset: 0
+        staircase: true
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 0

official/projects/qat/vision/configs/experiments/image_classification/imagenet_mobilenetv3large_qat_tpu.yaml

+runtime:
+  distribution_strategy: 'tpu'
+  mixed_precision_dtype: 'float32'
+task:
+  model:
+    num_classes: 1001
+    input_size: [224, 224, 3]
+    backbone:
+      type: 'mobilenet'
+      mobilenet:
+        model_id: 'MobileNetV3Large'
+        filter_size_scale: 1.0
+    dropout_rate: 0.3
+  losses:
+    l2_weight_decay: 1.0e-06  # 1/10 of original value.
+    one_hot: true
+    label_smoothing: 0.1
+  train_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/train*'
+    is_training: true
+    global_batch_size: 4096
+    dtype: 'float32'
+    aug_rand_hflip: true
+    aug_type:
+      autoaug:
+        augmentation_name: v0
+        cutout_const: 100
+        translate_const: 250
+      type: autoaug
+    drop_remainder: true
+  validation_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/valid*'
+    is_training: false
+    global_batch_size: 4096
+    dtype: 'float32'
+    drop_remainder: false
+    aug_rand_hflip: true
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/mobilenetv3_baseline_31/ckpt-156000'
+trainer:
+  # With below setting, the accuracy of QAT reaches to accuracy 0.74.43 after ~2 hours with 4x4 DF.
+  train_steps: 62400
+  validation_steps: 13
+  validation_interval: 312
+  steps_per_loop: 312
+  summary_interval: 312
+  checkpoint_interval: 312
+  optimizer_config:
+    learning_rate:
+      cosine:
+        alpha: 0.0
+        decay_steps: 62400
+        initial_learning_rate: 0.0003  # 1/10 of original lr.
+        name: CosineDecay
+        offset: 0
+      type: cosine
+    optimizer:
+      adamw:
+        amsgrad: false
+        beta_1: 0.9
+        beta_2: 0.999
+        epsilon: 1.0e-07
+        gradient_clip_norm: 1.0
+        weight_decay_rate: 0.0
+      type: adamw
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 0

official/projects/qat/vision/configs/experiments/image_classification/imagenet_resnet50_qat_gpu.yaml

+runtime:
+  distribution_strategy: 'mirrored'
+  mixed_precision_dtype: 'float32'
+  loss_scale: 'dynamic'
+task:
+  model:
+    num_classes: 1001
+    input_size: [224, 224, 3]
+    backbone:
+      type: 'resnet'
+      resnet:
+        model_id: 50
+  losses:
+    l2_weight_decay: 0.0001
+    one_hot: true
+    label_smoothing: 0.1
+  train_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/train*'
+    is_training: true
+    global_batch_size: 256
+    dtype: 'float32'
+  validation_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/valid*'
+    is_training: false
+    global_batch_size: 256
+    dtype: 'float32'
+    drop_remainder: false
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/resnet_classifier_gpu/ckpt-56160'
+trainer:
+  # With below setting, the accuracy of QAT reaches to Top1-accuracy 0.7720 after 5 days of training
+  # with 8GPUs, which is higher than the non-quantized float32 version Resnet.
+  train_steps: 449280
+  validation_steps: 200
+  validation_interval: 5000
+  steps_per_loop: 5000
+  summary_interval: 5000
+  checkpoint_interval: 5000
+  optimizer_config:
+    optimizer:
+      type: 'sgd'
+      sgd:
+        momentum: 0.9
+    learning_rate:
+      type: 'stepwise'
+      stepwise:
+        boundaries: [150000, 300000, 400000]
+        values: [0.08, 0.008, 0.0008, 0.00008]
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 40000

official/projects/qat/vision/configs/experiments/image_classification/imagenet_resnet50_qat_gpu_fast.yaml

+runtime:
+  distribution_strategy: 'mirrored'
+  mixed_precision_dtype: 'float32'
+  loss_scale: 'dynamic'
+task:
+  model:
+    num_classes: 1001
+    input_size: [224, 224, 3]
+    backbone:
+      type: 'resnet'
+      resnet:
+        model_id: 50
+  losses:
+    l2_weight_decay: 0.0001
+    one_hot: true
+    label_smoothing: 0.1
+  train_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/train*'
+    is_training: true
+    global_batch_size: 256
+    dtype: 'float32'
+  validation_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/valid*'
+    is_training: false
+    global_batch_size: 256
+    dtype: 'float32'
+    drop_remainder: false
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/resnet_classifier_gpu/ckpt-56160'
+trainer:
+  # With below setting, the accuracy of QAT reaches to the non-quantized float32 version after
+  # around 160k steps, which takes 1d 15h with 8 GPUS.
+  train_steps: 449280
+  validation_steps: 200
+  validation_interval: 5000
+  steps_per_loop: 5000
+  summary_interval: 5000
+  checkpoint_interval: 5000
+  optimizer_config:
+    optimizer:
+      type: 'sgd'
+      sgd:
+        momentum: 0.9
+    learning_rate:
+      type: 'exponential'
+      exponential:
+        initial_learning_rate: 0.016
+        decay_steps: 25000
+        decay_rate: 0.5
+        staircase: true
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 1000

official/projects/qat/vision/configs/experiments/image_classification/imagenet_resnet50_qat_gpu_fast_4x4.yaml

+runtime:
+  distribution_strategy: 'mirrored'
+  mixed_precision_dtype: 'float32'
+  loss_scale: 'dynamic'
+task:
+  model:
+    num_classes: 1001
+    input_size: [224, 224, 3]
+    backbone:
+      type: 'resnet'
+      resnet:
+        model_id: 50
+  losses:
+    l2_weight_decay: 0.0001
+    one_hot: true
+    label_smoothing: 0.1
+  train_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/train*'
+    is_training: true
+    global_batch_size: 256
+    dtype: 'float32'
+  validation_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/valid*'
+    is_training: false
+    global_batch_size: 256
+    dtype: 'float32'
+    drop_remainder: false
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/resnet_classifier_gpu/ckpt-56160'
+    change_num_bits: true
+    num_bits_weight: 4
+    num_bits_activation: 4
+trainer:
+  # With below setting, the accuracy of QAT reaches Top1-accuracy 0.6822 at 205k steps with 8GPUs.
+  # TODO: Please change the configs when training is done.
+  train_steps: 449280
+  validation_steps: 200
+  validation_interval: 5000
+  steps_per_loop: 5000
+  summary_interval: 5000
+  checkpoint_interval: 5000
+  optimizer_config:
+    optimizer:
+      type: 'sgd'
+      sgd:
+        momentum: 0.9
+    learning_rate:
+      type: 'exponential'
+      exponential:
+        initial_learning_rate: 0.016
+        decay_steps: 25000
+        decay_rate: 0.5
+        staircase: true
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 1000

official/projects/qat/vision/configs/experiments/image_classification/imagenet_resnet50_qat_gpu_fast_4x8.yaml

+runtime:
+  distribution_strategy: 'mirrored'
+  mixed_precision_dtype: 'float32'
+  loss_scale: 'dynamic'
+task:
+  model:
+    num_classes: 1001
+    input_size: [224, 224, 3]
+    backbone:
+      type: 'resnet'
+      resnet:
+        model_id: 50
+  losses:
+    l2_weight_decay: 0.0001
+    one_hot: true
+    label_smoothing: 0.1
+  train_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/train*'
+    is_training: true
+    global_batch_size: 256
+    dtype: 'float32'
+  validation_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/valid*'
+    is_training: false
+    global_batch_size: 256
+    dtype: 'float32'
+    drop_remainder: false
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/resnet_classifier_gpu/ckpt-56160'
+    change_num_bits: true
+    num_bits_weight: 4
+    num_bits_activation: 8
+trainer:
+  # With below setting, the accuracy of QAT reaches Top1-accuracy 0.7575 at 220k steps with 8GPUs.
+  # TODO: Please change the configs when training is done.
+  train_steps: 449280
+  validation_steps: 200
+  validation_interval: 5000
+  steps_per_loop: 5000
+  summary_interval: 5000
+  checkpoint_interval: 5000
+  optimizer_config:
+    optimizer:
+      type: 'sgd'
+      sgd:
+        momentum: 0.9
+    learning_rate:
+      type: 'exponential'
+      exponential:
+        initial_learning_rate: 0.016
+        decay_steps: 25000
+        decay_rate: 0.5
+        staircase: true
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 1000

official/projects/qat/vision/configs/experiments/image_classification/imagenet_resnet50_qat_gpu_fast_6x6.yaml

+runtime:
+  distribution_strategy: 'mirrored'
+  mixed_precision_dtype: 'float32'
+  loss_scale: 'dynamic'
+task:
+  model:
+    num_classes: 1001
+    input_size: [224, 224, 3]
+    backbone:
+      type: 'resnet'
+      resnet:
+        model_id: 50
+  losses:
+    l2_weight_decay: 0.0001
+    one_hot: true
+    label_smoothing: 0.1
+  train_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/train*'
+    is_training: true
+    global_batch_size: 256
+    dtype: 'float32'
+  validation_data:
+    input_path: '/readahead/200M/placer/prod/home/distbelief/imagenet-tensorflow/imagenet-2012-tfrecord/valid*'
+    is_training: false
+    global_batch_size: 256
+    dtype: 'float32'
+    drop_remainder: false
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/resnet_classifier_gpu/ckpt-56160'
+    change_num_bits: true
+    num_bits_weight: 6
+    num_bits_activation: 6
+trainer:
+  # With below setting, the accuracy of QAT reaches Top1-accuracy 0.7607 at 190k steps with 8GPUs.
+  # TODO: Please change the configs when training is done.
+  train_steps: 449280
+  validation_steps: 200
+  validation_interval: 5000
+  steps_per_loop: 5000
+  summary_interval: 5000
+  checkpoint_interval: 5000
+  optimizer_config:
+    optimizer:
+      type: 'sgd'
+      sgd:
+        momentum: 0.9
+    learning_rate:
+      type: 'exponential'
+      exponential:
+        initial_learning_rate: 0.016
+        decay_steps: 25000
+        decay_rate: 0.5
+        staircase: true
+    warmup:
+      type: 'linear'
+      linear:
+        warmup_steps: 1000

official/projects/qat/vision/configs/experiments/retinanet/coco_spinenet49_mobile_qat_gpu.yaml

+# --experiment_type=retinanet_spinenet_mobile_coco_qat
+runtime:
+  distribution_strategy: 'mirrored'
+  mixed_precision_dtype: 'float32'
+task:
+  losses:
+    l2_weight_decay: 3.0e-05
+  model:
+    anchor:
+      anchor_size: 3
+      aspect_ratios: [0.5, 1.0, 2.0]
+      num_scales: 3
+    backbone:
+      spinenet_mobile:
+        stochastic_depth_drop_rate: 0.2
+        model_id: '49'
+        se_ratio: 0.2
+        use_keras_upsampling_2d: true
+      type: 'spinenet_mobile'
+    decoder:
+      type: 'identity'
+    head:
+      num_convs: 4
+      num_filters: 48
+      use_separable_conv: true
+    input_size: [384, 384, 3]
+    max_level: 7
+    min_level: 3
+    norm_activation:
+      activation: 'swish'
+      norm_epsilon: 0.001
+      norm_momentum: 0.99
+      use_sync_bn: true
+  train_data:
+    dtype: 'float32'
+    global_batch_size: 128
+    is_training: true
+    parser:
+      aug_rand_hflip: true
+      aug_scale_max: 2.0
+      aug_scale_min: 0.5
+  validation_data:
+    dtype: 'float32'
+    global_batch_size: 8
+    is_training: false
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/coco_spinenet49_mobile_tpu/ckpt-277200'
+trainer:
+  checkpoint_interval: 924
+  optimizer_config:
+    learning_rate:
+      stepwise:
+        boundaries: [531300, 545160]
+        values: [0.0016, 0.00016, 0.000016]
+      type: 'stepwise'
+    warmup:
+      linear:
+        warmup_learning_rate: 0.0000335
+        warmup_steps: 4000
+  steps_per_loop: 924
+  train_steps: 554400
+  validation_interval: 924
+  validation_steps: 1250
+  summary_interval: 924

official/projects/qat/vision/configs/experiments/semantic_segmentation/deeplabv3_mobilenetv2_pascal_qat_gpu.yaml

+# --experiment_type=mnv2_deeplabv3_pascal_qat
+# Use 8 v100 GPUs for training and 4 v100 GPUs for eval.
+# mIoU (unquantized fp32): 74.78
+runtime:
+  distribution_strategy: 'mirrored'
+  mixed_precision_dtype: 'float32'
+  loss_scale: 'dynamic'
+task:
+  model:
+    num_classes: 21
+    input_size: [512, 512, 3]
+    backbone:
+      type: 'mobilenet'
+      mobilenet:
+        model_id: 'MobileNetV2'
+        output_stride: 16
+    decoder:
+      aspp:
+        dilation_rates: []
+        level: 4
+        pool_kernel_size: null
+        output_tensor: true
+      type: 'aspp'
+    head:
+      feature_fusion: null
+      num_convs: 0
+    norm_activation:
+      activation: relu
+      norm_epsilon: 0.001
+      norm_momentum: 0.99
+      use_sync_bn: true
+  losses:
+    l2_weight_decay: 4.0e-07  # 1/100 of original value.
+  train_data:
+    output_size: [512, 512]
+    crop_size: [512, 512]
+    input_path: 'gs://**/pascal_voc_seg/train_aug*'
+    is_training: true
+    global_batch_size: 16
+    dtype: 'float32'
+    aug_rand_hflip: true
+    aug_scale_max: 2.0
+    aug_scale_min: 0.5
+  validation_data:
+    output_size: [512, 512]
+    input_path: 'gs://**/pascal_voc_seg/val*'
+    is_training: false
+    global_batch_size: 16
+    dtype: 'float32'
+    drop_remainder: false
+    resize_eval_groundtruth: false
+    groundtruth_padded_size: [512, 512]
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/deeplabv3_mobilenetv2_pascal_coco_0.21/29808901/best_ckpt/best_ckpt-54'
+  init_checkpoint: null
+trainer:
+  optimizer_config:
+    learning_rate:
+      polynomial:
+        decay_steps: 13240
+        initial_learning_rate: 0.00007  # 1/100 of original lr.
+        power: 0.9
+      type: polynomial
+    optimizer:
+      sgd:
+        momentum: 0.9
+      type: sgd
+    warmup:
+      linear:
+        name: linear
+        warmup_steps: 0  # No warmup
+      type: linear
+  best_checkpoint_eval_metric: 'mean_iou'
+  best_checkpoint_export_subdir: 'best_ckpt'
+  best_checkpoint_metric_comp: 'higher'
+  steps_per_loop: 662
+  summary_interval: 662
+  train_steps: 13240
+  validation_interval: 662
+  validation_steps: 90
+  checkpoint_interval: 662

official/projects/qat/vision/configs/experiments/semantic_segmentation/deeplabv3_mobilenetv2_pascal_qat_tpu.yaml

+# --experiment_type=mnv2_deeplabv3_pascal_qat
+# Use 4x2 DF for training and eval.
+# mIoU (unquantized fp32): 74.69
+runtime:
+  distribution_strategy: 'tpu'
+  mixed_precision_dtype: 'float32'
+task:
+  model:
+    num_classes: 21
+    input_size: [512, 512, 3]
+    backbone:
+      type: 'mobilenet'
+      mobilenet:
+        model_id: 'MobileNetV2'
+        output_stride: 16
+    decoder:
+      aspp:
+        dilation_rates: []
+        level: 4
+        pool_kernel_size: null
+        output_tensor: true
+      type: 'aspp'
+    head:
+      feature_fusion: null
+      num_convs: 0
+    norm_activation:
+      activation: relu
+      norm_epsilon: 0.001
+      norm_momentum: 0.99
+      use_sync_bn: true
+  losses:
+    l2_weight_decay: 4.0e-07  # 1/100 of original value.
+  train_data:
+    output_size: [512, 512]
+    crop_size: [512, 512]
+    input_path: 'gs://**/pascal_voc_seg/train_aug*'
+    is_training: true
+    global_batch_size: 16
+    dtype: 'float32'
+    aug_rand_hflip: true
+    aug_scale_max: 2.0
+    aug_scale_min: 0.5
+  validation_data:
+    output_size: [512, 512]
+    input_path: 'gs://**/pascal_voc_seg/val*'
+    is_training: false
+    global_batch_size: 16
+    dtype: 'float32'
+    drop_remainder: false
+    resize_eval_groundtruth: false
+    groundtruth_padded_size: [512, 512]
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/deeplabv3_mobilenetv2_pascal_coco_0.21/29808901/best_ckpt/best_ckpt-54'
+  init_checkpoint: null
+trainer:
+  optimizer_config:
+    learning_rate:
+      polynomial:
+        decay_steps: 13240
+        initial_learning_rate: 0.00007  # 1/100 of original lr.
+        power: 0.9
+      type: polynomial
+    optimizer:
+      sgd:
+        momentum: 0.9
+      type: sgd
+    warmup:
+      linear:
+        name: linear
+        warmup_steps: 0  # No warmup
+      type: linear
+  best_checkpoint_eval_metric: 'mean_iou'
+  best_checkpoint_export_subdir: 'best_ckpt'
+  best_checkpoint_metric_comp: 'higher'
+  steps_per_loop: 662
+  summary_interval: 662
+  train_steps: 13240
+  validation_interval: 662
+  validation_steps: 90
+  checkpoint_interval: 662

official/projects/qat/vision/configs/experiments/semantic_segmentation/deeplabv3plus_mobilenetv2_cityscapes_qat_tpu.yaml

+# --experiment_type=mnv2_deeplabv3plus_cityscapes_qat
+# Use 4x2 DF for training and eval.
+# mIoU (unquantized fp32): 73.84
+runtime:
+  distribution_strategy: 'tpu'
+  mixed_precision_dtype: 'float32'
+task:
+  model:
+    num_classes: 19
+    input_size: [1024, 2048, 3]
+    backbone:
+      type: 'mobilenet'
+      mobilenet:
+        model_id: 'MobileNetV2'
+        output_stride: 16
+        output_intermediate_endpoints: true
+    decoder:
+      aspp:
+        dilation_rates: []
+        level: 4
+        pool_kernel_size: [512, 1024]
+        output_tensor: true
+      type: 'aspp'
+    head:
+      feature_fusion: 'deeplabv3plus'
+      low_level: '2/depthwise'
+      low_level_num_filters: 48
+      level: 4
+      num_convs: 2
+      use_depthwise_convolution: true
+    norm_activation:
+      activation: relu
+      norm_epsilon: 0.001
+      norm_momentum: 0.99
+      use_sync_bn: true
+  losses:
+    l2_weight_decay: 4.0e-07  # 1/100 of original value.
+  train_data:
+    output_size: [1024, 2048]
+    crop_size: []
+    input_path: ''
+    tfds_name: 'cityscapes/semantic_segmentation'
+    tfds_split: 'train'
+    is_training: true
+    global_batch_size: 16
+    dtype: 'float32'
+    aug_rand_hflip: true
+    aug_scale_max: 2.0
+    aug_scale_min: 0.5
+  validation_data:
+    output_size: [1024, 2048]
+    input_path: ''
+    tfds_name: 'cityscapes/semantic_segmentation'
+    tfds_split: 'validation'
+    is_training: false
+    global_batch_size: 16
+    dtype: 'float32'
+    drop_remainder: false
+    resize_eval_groundtruth: true
+  quantization:
+    pretrained_original_checkpoint: 'gs://**/deeplabv3plus_mobilenetv2_cityscapes/29814723/best_ckpt/best_ckpt-408'
+  init_checkpoint: null
+trainer:
+  optimizer_config:
+    learning_rate:
+      polynomial:
+        decay_steps: 20000
+        initial_learning_rate: 0.0001  # 1/100 of original lr.
+        power: 0.9
+      type: polynomial
+    optimizer:
+      sgd:
+        momentum: 0.9
+      type: sgd
+    warmup:
+      linear:
+        name: linear
+        warmup_learning_rate: 0
+        warmup_steps: 0  # No warmup
+      type: linear
+  steps_per_loop: 185
+  summary_interval: 185
+  train_steps: 20000
+  validation_interval: 185
+  validation_steps: 31
+  checkpoint_interval: 185
+  best_checkpoint_export_subdir: 'best_ckpt'
+  best_checkpoint_eval_metric: 'mean_iou'
+  best_checkpoint_metric_comp: 'higher'

official/projects/qat/vision/configs/image_classification.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.
+
+# Lint as: python3
+"""Image classification configuration definition."""
+
+import dataclasses
+from typing import Optional
+
+from official.core import config_definitions as cfg
+from official.core import exp_factory
+from official.projects.qat.vision.configs import common
+from official.vision.beta.configs import image_classification
+
+
+@dataclasses.dataclass
+class ImageClassificationTask(image_classification.ImageClassificationTask):
+  quantization: Optional[common.Quantization] = None
+
+
+@exp_factory.register_config_factory('resnet_imagenet_qat')
+def image_classification_imagenet() -> cfg.ExperimentConfig:
+  """Builds an image classification config for the resnet with QAT."""
+  config = image_classification.image_classification_imagenet()
+  task = ImageClassificationTask.from_args(
+      quantization=common.Quantization(), **config.task.as_dict())
+  config.task = task
+
+  return config
+
+
+@exp_factory.register_config_factory('mobilenet_imagenet_qat')
+def image_classification_imagenet_mobilenet() -> cfg.ExperimentConfig:
+  """Builds an image classification config for the mobilenetV2 with QAT."""
+  config = image_classification.image_classification_imagenet_mobilenet()
+  task = ImageClassificationTask.from_args(
+      quantization=common.Quantization(), **config.task.as_dict())
+  config.task = task
+
+  return config

official/projects/qat/vision/configs/image_classification_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 for image_classification."""
+# pylint: disable=unused-import
+from absl.testing import parameterized
+import tensorflow as tf
+
+from official.core import config_definitions as cfg
+from official.core import exp_factory
+from official.projects.qat.vision.configs import common
+from official.projects.qat.vision.configs import image_classification as qat_exp_cfg
+from official.vision import beta
+from official.vision.beta.configs import image_classification as exp_cfg
+
+
+class ImageClassificationConfigTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      ('resnet_imagenet_qat',),
+      ('mobilenet_imagenet_qat',),
+  )
+  def test_image_classification_configs(self, config_name):
+    config = exp_factory.get_exp_config(config_name)
+    self.assertIsInstance(config, cfg.ExperimentConfig)
+    self.assertIsInstance(config.task, qat_exp_cfg.ImageClassificationTask)
+    self.assertIsInstance(config.task.model,
+                          exp_cfg.ImageClassificationModel)
+    self.assertIsInstance(config.task.quantization, common.Quantization)
+    self.assertIsInstance(config.task.train_data, exp_cfg.DataConfig)
+    config.task.train_data.is_training = None
+    with self.assertRaisesRegex(KeyError, 'Found inconsistncy between key'):
+      config.validate()
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/projects/qat/vision/configs/retinanet.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.
+
+# Lint as: python3
+"""RetinaNet configuration definition."""
+import dataclasses
+from typing import Optional
+
+from official.core import config_definitions as cfg
+from official.core import exp_factory
+from official.projects.qat.vision.configs import common
+from official.vision.beta.configs import retinanet
+from official.vision.beta.configs.google import backbones
+
+
+@dataclasses.dataclass
+class RetinaNetTask(retinanet.RetinaNetTask):
+  quantization: Optional[common.Quantization] = None
+
+
+@exp_factory.register_config_factory('retinanet_spinenet_mobile_coco_qat')
+def retinanet_spinenet_mobile_coco() -> cfg.ExperimentConfig:
+  """Generates a config for COCO OD RetinaNet for mobile with QAT."""
+  config = retinanet.retinanet_spinenet_mobile_coco()
+  task = RetinaNetTask.from_args(
+      quantization=common.Quantization(), **config.task.as_dict())
+  task.model.backbone = backbones.Backbone(
+      type='spinenet_mobile',
+      spinenet_mobile=backbones.SpineNetMobile(
+          model_id='49',
+          stochastic_depth_drop_rate=0.2,
+          min_level=3,
+          max_level=7,
+          use_keras_upsampling_2d=True))
+  config.task = task
+
+  return config