MobileDet CPU doesn't support QAT training

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.

1. The entire URL of the file you are using

https://github.com/tensorflow/models/tree/master/research/object_detection

2. Describe the bug

MobileDet CPU doesn't support QAT training.

3. Steps to reproduce

Just add QAT in Coco training file (ssdlite_mobiledet_cpu_320x320_coco_sync_4x4.config)

graph_rewriter {
  quantization {
    delay: 50
    activation_bits: 8
    weight_bits: 8
  }

Then, it fails after saving the first training checkpoint after starting doing the quantization. Exactly, when it tries to restore the model from the checkpoint.

tensorflow.python.framework.errors_impl.NotFoundError: Restoring from checkpoint failed. This is most likely due to a Variable name or other graph key that is missing from the checkpoint. Please ensure that you have not altered the graph expected based on the checkpoint. Original error:

Key FeatureExtractor/MobileDetCPU/Conv/conv_quant/max not found in checkpoint
	 [[node save/RestoreV2 (defined at /tensorflow-1.15.2/python3.6/tensorflow_core/python/framework/ops.py:1748) ]]

4. Expected behavior

It should work.

5. Additional context

Also, this checkpoint fails if you try to convert to inference graph with the python object_detection/export_inference_graph.py or python object_detection/export_tflite_ssd_graph.py command. It fails with the same error:

tensorflow.python.framework.errors_impl.NotFoundError: Key FeatureExtractor/MobileDetCPU/Conv/conv_quant/max not found in checkpoint
	 [[node save/RestoreV2 (defined at /Users/vferrer/miniconda3/envs/Pytorch/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1748) ]]

Neither works disable checkpoints, performing the full training. It fails with the same error in python object_detection/export_inference_graph.py or python object_detection/export_tflite_ssd_graph.py and at the end of the training, when it tries to restore the model from the checkpoint.

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): colab
• TensorFlow version (use command below): 1.15.3
• Python version: colab
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version:
• GPU model and memory: k80