Node name is not unique when freezing axial-deep lab model from deeplab2

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 been filed already.

1. The entire URL of the file you are using

https://storage.googleapis.com/gresearch/tf-deeplab/checkpoint/max_deeplab_l_backbone_os16_axial_deeplab_cityscapes_trainfine.tar.gz

2. Describe the bug

From google/deeplab2/export_model.py, add the following line:

  signatures = module.__call__.get_concrete_function(module.get_input_spec())
  #following is the added code.
  frozen_func = convert_variables_to_constants_v2(signatures)

full command used

 python export.py
--experiment_option_path
"./configs/cityscapes/axial_deeplab/max_deeplab_l_backbone_os16.textproto"
--checkpoint_path
"C:/develop/max_deeplab_l_backbone_os16_axial_deeplab_cityscapes_trainfine_saved_model/ckpt-60000"

The following error pops up:

ValueError: Node 'DeepLab/max_deeplab_l_backbone/stage4/block2/attention/height_axis/query_rpe/Gather/axis' is not unique
A clear and concise description of what the bug is.

3. Steps to reproduce

A quick way to reproduce is , after downloaded the checkpoint, use the code below:

model = tf.keras.models.load_model(r"C:/develop/max_deeplab_l_backbone_os16_axial_deeplab_cityscapes_trainfine_saved_model/")
 signatures = model.signatures['serving_default']
 frozen_func = convert_variables_to_constants_v2(signatures)

4. Expected behavior

Able to generate a frozen graph that can be exported.

5. Additional context

2022-01-21 17:28:22.472844: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'cudnn64_8.dll'; dlerror: cudnn64_8.dll not found
2022-01-21 17:28:22.473118: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1835] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
2022-01-21 17:28:33.104143: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: AVX AVX2
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
I0121 17:28:33.102456 21356 deeplab.py:57] Synchronized Batchnorm is used.
I0121 17:28:33.110966 21356 axial_resnet_instances.py:144] Axial-ResNet final config: {'num_blocks': [3, 6, 3, 3], 'backbone_layer_multiplier': 1.0, 'width_multiplier': 1.0, 'stem_width_multiplier': 1.0, 'output_stride': 16, 'classification_mode': True, 'backbone_type': 'wider_resnet', 'use_axial_beyond_stride': 16, 'backbone_use_transformer_beyond_stride': 0, 'extra_decoder_use_transformer_beyond_stride': 16, 'backbone_decoder_num_stacks': 1, 'backbone_decoder_blocks_per_stage': 1, 'extra_decoder_num_stacks': 1, 'extra_decoder_blocks_per_stage': 3, 'max_num_mask_slots': 128, 'num_mask_slots': 128, 'memory_channels': 512, 'base_transformer_expansion': 2.0, 'global_feed_forward_network_channels': 512, 'high_resolution_output_stride': 4, 'activation': 'relu', 'block_group_config': {'attention_bottleneck_expansion': 4, 'drop_path_keep_prob': 1.0, 'drop_path_beyond_stride': 4, 'drop_path_schedule': 'linear', 'positional_encoding_type': None, 'use_global_beyond_stride': 0, 'use_sac_beyond_stride': -1, 'use_squeeze_and_excite': False, 'conv_use_recompute_grad': False, 'axial_use_recompute_grad': True, 'recompute_within_stride': 0, 'transformer_use_recompute_grad': False, 'axial_layer_config': {'query_shape': (129, 129), 'key_expansion': 2, 'value_expansion': 4, 'memory_flange': (32, 32), 'double_global_attention': False, 'num_heads': 8, 'use_query_rpe_similarity': True, 'use_key_rpe_similarity': True, 'use_content_similarity': True, 'retrieve_value_rpe': True, 'retrieve_value_content': True, 'initialization_std_for_query_key_rpe': 1.0, 'initialization_std_for_value_rpe': 1.0, 'self_attention_activation': 'softmax'}, 'dual_path_transformer_layer_config': {'num_heads': 8, 'bottleneck_expansion': 2, 'key_expansion': 1, 'value_expansion': 2, 'feed_forward_network_channels': 2048, 'use_memory_self_attention': True, 'use_pixel2memory_feedback_attention': True, 'transformer_activation': 'softmax'}}, 'bn_layer': functools.partial(<class 'keras.layers.normalization.batch_normalization.SyncBatchNormalization'>, momentum=0.9900000095367432, epsilon=0.0010000000474974513), 'conv_kernel_weight_decay': 0.0}
I0121 17:28:33.613099 21356 deeplab.py:96] Setting pooling size to (65, 129)
I0121 17:28:36.232671 21356 api.py:446] Eval with scales ListWrapper([1.0])
I0121 17:28:39.102459 21356 api.py:446] Eval scale 1.0; setting pooling size to [None, None]
I0121 17:28:58.836912 21356 api.py:446] Eval with scales ListWrapper([1.0])
I0121 17:28:58.883746 21356 api.py:446] Eval scale 1.0; setting pooling size to [None, None]
2022-01-21 17:40:38.838944: I tensorflow/core/grappler/devices.cc:66] Number of eligible GPUs (core count >= 8, compute capability >= 0.0): 1
2022-01-21 17:40:38.847874: I tensorflow/core/grappler/clusters/single_machine.cc:357] Starting new session
2022-01-21 17:40:38.883561: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1835] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
2022-01-21 17:40:40.300834: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:1137] Optimization results for grappler item: graph_to_optimize
function_optimizer: Graph size after: 4654 nodes (1206), 6100 edges (1886), time = 239.541ms.
function_optimizer: function_optimizer did nothing. time = 2.73ms.
Traceback (most recent call last):
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\tensorflow\python\framework\importer.py", line 497, in _import_graph_def_internal
graph._c_graph, serialized, options) # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: Node 'DeepLab/max_deeplab_l_backbone/stage4/block2/attention/height_axis/query_rpe/Gather/axis' is not unique
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\absl\app.py", line 303, in run
_run_main(main, args)
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\absl\app.py", line 251, in _run_main
sys.exit(main(argv))
File "C:/develop/deeplab2/export_model.py", line 158, in main
frozen_func = convert_variables_to_constants_v2(signatures)
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\tensorflow\python\framework\convert_to_constants.py", line 1154, in convert_variables_to_constants_v2
converted_input_indices)
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\tensorflow\python\framework\convert_to_constants.py", line 1080, in _construct_concrete_function
new_output_names)
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\tensorflow\python\eager\wrap_function.py", line 650, in function_from_graph_def
wrapped_import = wrap_function(_imports_graph_def, [])
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\tensorflow\python\eager\wrap_function.py", line 628, in wrap_function
collections={}),
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\tensorflow\python\framework\func_graph.py", line 1007, in func_graph_from_py_func
func_outputs = python_func(*func_args, **func_kwargs)
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\tensorflow\python\eager\wrap_function.py", line 87, in call
return self.call_with_variable_creator_scope(self._fn)(*args, **kwargs)
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\tensorflow\python\eager\wrap_function.py", line 93, in wrapped
return fn(*args, **kwargs)
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\tensorflow\python\eager\wrap_function.py", line 648, in _imports_graph_def
importer.import_graph_def(graph_def, name="")
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\tensorflow\python\util\deprecation.py", line 549, in new_func
return func(*args, **kwargs)
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\tensorflow\python\framework\importer.py", line 405, in import_graph_def
producer_op_list=producer_op_list)
File "C:\Users\tfan.conda\envs\XNNC\lib\site-packages\tensorflow\python\framework\importer.py", line 501, in _import_graph_def_internal
raise ValueError(str(e))
ValueError: Node 'DeepLab/max_deeplab_l_backbone/stage4/block2/attention/height_axis/query_rpe/Gather/axis' is not unique

6. System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 10
• Mobile device name if the issue happens on a mobile device:
• TensorFlow installed from (source or binary): pip
• TensorFlow version (use command below):tried 2.5/2.7/nightly
• Python version: 3.6.9
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version: NA
• GPU model and memory: not using gpu