Add support for QAT testing using mobilenet_trainer.py

Prerequisites

Please answer the following questions for yourself before submitting an issue.

• I am using the latest TensorFlow Model Garden release and TensorFlow 2.
• I am reporting the issue to the correct repository. (Model Garden official or research directory)
• I checked to make sure that this issue has not already been filed.

I am using the latest mobilenet branch in TensorFlow Models repo and TensorFlow 2.3.0.

1. The entire URL of the file you are using

https://github.com/tensorflow/models/blob/mobilenet/research/mobilenet/mobilenet_trainer.py

2. Describe the bug

Currently, the mobilenet_trainer file does not support loading pre-trained mobilenet_v1 from a TF 1.x checkpoint here. If we try to load the checkpoint file via flag resume_checkpoint, it throws an error AssertionError: Nothing except the root object matched a checkpointed value. Typically this means that the checkpoint does not match the Python program.

3. Steps to reproduce

To reproduce the error, we can run:
python -m research.mobilenet.mobilenet_trainer --model_name=mobilenet_v1 --dataset_name=imagenet2012 --data_dir=/mnt/ssd/imagenet/tfrecords_official_models/imagenet --resume_checkpoint=True --model_dir=/work/forks/models/research/mobilenet/end_to_end_test/mobilenet_v1_ck

4. Expected behavior

A clear and concise description of what you expected to happen.

5. Additional context

There are also other issues: for instance, relu6 caused issues when the QAT API from tfmot is applied, the wrong num_of_classes which causes accuracy drop, and the batch_size for training is not optimal. Together with the checkpoint loading, these have been all fixed locally (the baseline evaluation gives Top-1 accuracy of 0.71 and the QAT evaluation gives 0.708), but need to see whether I can push the PR.

6. System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
• Mobile device name if the issue happens on a mobile device:
• TensorFlow installed from (source or binary):
• TensorFlow version (use command below):
• Python version:
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version:
• GPU model and memory: