Custom Training with TF 2 Object Detect API Fails

Prerequisites

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

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

1. The entire URL of the file you are using

https://github.com/tensorflow/models/tree/master/official/...

2. Describe the bug

I'm attempting to migrate from the TF1 object detect API to the TF2 object detect API. The exact model isn't available in the TF2 version (quantized SSD Mobilenet v2), so I'm using EfficientDet-d0. But I'm attempting to keep as many things the same as possible, including using the exact same tfrecord training and validation files, similar config settings, etc. And I'm starting the fine-tuning from the config + weights found in the TF zoo. Note that there are 4 classes, and the config and label_map files have been updated appropriately.

This is being run locally on a machine with 2 GPUs, and I had to use tf.config.experimental.set_memory_growth(gpu, True) to get it to run at all.

Two failure modes are observed:

1. num_workers = 1: prints a long list of warnings regarding unresolved objects in the checkpoint, then exits without printing any errors. Here's the tail of the console output:

WARNING:tensorflow:Unresolved object in checkpoint: (root).model._feature_extractor._bifpn_stage.node_input_blocks.7.0.1.1.moving_mean
W0723 12:41:23.119486 139840608290624 util.py:143] Unresolved object in checkpoint: (root).model._feature_extractor._bifpn_stage.node_input_blocks.7.0.1.1.moving_mean
WARNING:tensorflow:Unresolved object in checkpoint: (root).model._feature_extractor._bifpn_stage.node_input_blocks.7.0.1.1.moving_variance
W0723 12:41:23.119561 139840608290624 util.py:143] Unresolved object in checkpoint: (root).model._feature_extractor._bifpn_stage.node_input_blocks.7.0.1.1.moving_variance
WARNING:tensorflow:A checkpoint was restored (e.g. tf.train.Checkpoint.restore or tf.keras.Model.load_weights) but not all checkpointed values were used. See above for specific issues. Use expect_partial() on the load status object, e.g. tf.train.Checkpoint.restore(...).expect_partial(), to silence these warnings, or use assert_consumed() to make the check explicit. See https://www.tensorflow.org/guide/checkpoint#loading_mechanics for details.
W0723 12:41:23.119644 139840608290624 util.py:151] A checkpoint was restored (e.g. tf.train.Checkpoint.restore or tf.keras.Model.load_weights) but not all checkpointed values were used. See above for specific issues. Use expect_partial() on the load status object, e.g. tf.train.Checkpoint.restore(...).expect_partial(), to silence these warnings, or use assert_consumed() to make the check explicit. See https://www.tensorflow.org/guide/checkpoint#loading_mechanics for details.

2. num_workers = 2: crashes with the following attribute error:

File "/home/user1/tf2odapi/models/research/object_detection/model_lib_v2.py", line 549, in train_loop
load_fine_tune_checkpoint(detection_model,
File "/home/user1/tf2odapi/models/research/object_detection/model_lib_v2.py", line 357, in load_fine_tune_checkpoint
strategy.run(
File "/home/user1/anaconda3/envs/tf2odapi/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py", line 951, in run
return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
File "/home/user1/anaconda3/envs/tf2odapi/lib/python3.8/site-packages/tensorflow/python/distribute/distribute_lib.py", line 2290, in call_for_each_replica
return self._call_for_each_replica(fn, args, kwargs)
File "/home/user1/anaconda3/envs/tf2odapi/lib/python3.8/site-packages/tensorflow/python/distribute/mirrored_strategy.py", line 743, in _call_for_each_replica
wrapped = self._cfer_fn_cache.get(fn)
AttributeError: 'CollectiveAllReduceExtended' object has no attribute '_cfer_fn_cache'

3. Steps to reproduce

Can't be reproduced exactly on your end, as it involves some local files.

4. Expected behavior

Expect training to run successfully.

5. Additional context

Include any logs that would be helpful to diagnose the problem.

6. System information

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