Error in loading OD checkpoints

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/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md
https://github.com/tensorflow/models/blob/master/research/object_detection/model_main_tf2.py

2. Describe the bug

model_main_tf2.py fails in assert_existing_objects_matched when importing fine_tune_checkpoint from models provided in the TF2 OD model zoo. If the assertion is suppressed the model is able to train but the checkpoint is not correctly loaded.

The error persists in all TF2 models tested:

• CenterNet HourGlass104 512x512
• EfficientDet D3 896x896
• SSD MobileNet v2 320x320

3. Steps to reproduce

Run model_main_tf2.py with fresh model zoo model.

4. Expected behavior

Loading checkpoint.

5. Additional context

fine_tune_checkpoint: "/code/ssd_mobilenet_v2_320x320_coco17_tpu-8/checkpoint/ckpt-0"

Traceback (most recent call last):
  File "/home/ubuntu/models/research/object_detection/model_main_tf2.py", line 113, in <module>
    tf.compat.v1.app.run()
  File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/platform/app.py", line 40, in run
    _run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
  File "/home/ubuntu/.local/lib/python3.6/site-packages/absl/app.py", line 299, in run
    _run_main(main, args)
  File "/home/ubuntu/.local/lib/python3.6/site-packages/absl/app.py", line 250, in _run_main
    sys.exit(main(argv))
  File "/home/ubuntu/models/research/object_detection/model_main_tf2.py", line 110, in main
    record_summaries=FLAGS.record_summaries)
  File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/object_detection/model_lib_v2.py", line 561, in train_loop
    unpad_groundtruth_tensors)
  File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/object_detection/model_lib_v2.py", line 375, in load_fine_tune_checkpoint
    ckpt.restore(checkpoint_path).assert_existing_objects_matched()
  File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow/python/training/tracking/util.py", line 783, in assert_existing_objects_matched
    (list(unused_python_objects),))
AssertionError: Some Python objects were not bound to checkpointed values, likely due to changes in the Python program:

Possibly irrelevant, but model_lib_v2.py still uses CheckpointV1 instead of the TF2 version defined in https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/training/tracking/util.py.

6. System information

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