Can't export Faster-RCNN with number_of_stages: 1

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/exporter_main_v2.py

2. Describe the bug

I tried to export the Faster R-CNN model with number_of_stages: 1 to get the RPN only and this is what I got

Traceback (most recent call last):
  File "exporter_main_v2.py", line 165, in <module>
    app.run(main)
  File "E:\Anaconda\envs\TFOD\lib\site-packages\absl\app.py", line 312, in run
    _run_main(main, args)
  File "E:\Anaconda\envs\TFOD\lib\site-packages\absl\app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "exporter_main_v2.py", line 158, in main
    exporter_lib_v2.export_inference_graph(
  File "E:\Anaconda\envs\TFOD\lib\site-packages\object_detection\exporter_lib_v2.py", line 271, in export_inference_graph
    concrete_function = detection_module.__call__.get_concrete_function()
  File "E:\Anaconda\envs\TFOD\lib\site-packages\tensorflow\python\eager\def_function.py", line 1299, in get_concrete_function
    concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
  File "E:\Anaconda\envs\TFOD\lib\site-packages\tensorflow\python\eager\def_function.py", line 1205, in _get_concrete_function_garbage_collected
    self._initialize(args, kwargs, add_initializers_to=initializers)
  File "E:\Anaconda\envs\TFOD\lib\site-packages\tensorflow\python\eager\def_function.py", line 725, in _initialize
    self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
  File "E:\Anaconda\envs\TFOD\lib\site-packages\tensorflow\python\eager\function.py", line 2969, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "E:\Anaconda\envs\TFOD\lib\site-packages\tensorflow\python\eager\function.py", line 3361, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "E:\Anaconda\envs\TFOD\lib\site-packages\tensorflow\python\eager\function.py", line 3196, in _create_graph_function
    func_graph_module.func_graph_from_py_func(
  File "E:\Anaconda\envs\TFOD\lib\site-packages\tensorflow\python\framework\func_graph.py", line 990, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "E:\Anaconda\envs\TFOD\lib\site-packages\tensorflow\python\eager\def_function.py", line 634, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "E:\Anaconda\envs\TFOD\lib\site-packages\tensorflow\python\framework\func_graph.py", line 977, in wrapper
    raise e.ag_error_metadata.to_exception(e)
tensorflow.python.autograph.pyct.error_utils.KeyError: in user code:

    E:\Anaconda\envs\TFOD\lib\site-packages\object_detection\exporter_lib_v2.py:163 call_func  *
        return self._run_inference_on_images(images, true_shapes, **kwargs)
    E:\Anaconda\envs\TFOD\lib\site-packages\object_detection\exporter_lib_v2.py:129 _run_inference_on_images  *
        detections[classes_field] = (

    KeyError: 'detection_classes'

3. Steps to reproduce

My pipeline.config file:

model {
  faster_rcnn {
    number_of_stages: 1 
    num_classes: 6
    image_resizer {
      keep_aspect_ratio_resizer {
        min_dimension: 300
        max_dimension: 500
        pad_to_max_dimension: true
      }
    }
    feature_extractor {
      type: "faster_rcnn_resnet50_keras"
    }
    first_stage_anchor_generator {
      grid_anchor_generator {
        height_stride: 4
        width_stride: 4
        scales: 0.25
        scales: 0.5
        scales: 1.0
        scales: 2.0
        aspect_ratios: 0.5
        aspect_ratios: 1.0
        aspect_ratios: 2.0
      }
    }
    first_stage_box_predictor_conv_hyperparams {
      op: CONV
      regularizer {
        l2_regularizer {
          weight: 0.0
        }
      }
      initializer {
        truncated_normal_initializer {
          stddev: 0.01
        }
      }
    }
    first_stage_nms_score_threshold: 0.0
    first_stage_nms_iou_threshold: 0.7
    first_stage_max_proposals: 300
    first_stage_localization_loss_weight: 2.0
    first_stage_objectness_loss_weight: 1.0
    initial_crop_size: 14
    maxpool_kernel_size: 2
    maxpool_stride: 2
    second_stage_box_predictor {
      mask_rcnn_box_predictor {
        fc_hyperparams {
          op: FC
          regularizer {
            l2_regularizer {
              weight: 0.0
            }
          }
          initializer {
            variance_scaling_initializer {
              factor: 1.0
              uniform: true
              mode: FAN_AVG
            }
          }
        }
        use_dropout: false
        dropout_keep_probability: 1.0
      }
    }
    second_stage_post_processing {
      batch_non_max_suppression {
        score_threshold: 0.0
        iou_threshold: 0.6
        max_detections_per_class: 100
        max_total_detections: 300
      }
      score_converter: SOFTMAX
    }
    second_stage_localization_loss_weight: 2.0
    second_stage_classification_loss_weight: 1.0
  }
}
train_config {
  batch_size: 16
  data_augmentation_options {
    random_adjust_hue {
    }
  }
  data_augmentation_options {
    random_adjust_contrast {
    }
  }
  data_augmentation_options {
    random_adjust_saturation {
    }
  }
  optimizer {
    momentum_optimizer {
      learning_rate {
        cosine_decay_learning_rate {
          learning_rate_base: 0.01
          total_steps: 200000
          warmup_learning_rate: 0.0
          warmup_steps: 5000
        }
      }
      momentum_optimizer_value: 0.9
    }
    use_moving_average: false
  }
  gradient_clipping_by_norm: 10.0
  fine_tune_checkpoint: "models/text_detection/faster_rcnn_resnet50_v1_800x1333/checkpoint/ckpt-0"
  num_steps: 200000
  fine_tune_checkpoint_type: "detection"
  fine_tune_checkpoint_version: V2
}
train_input_reader {
  label_map_path: "Data/CMT_Content/label_map.pbtxt"
  tf_record_input_reader {
    input_path: "Data/CMT_Content/Training/cmt_content_training.tfrecord"
  }
}
eval_config {
  metrics_set: "coco_detection_metrics"
  use_moving_averages: false
  batch_size: 1
}
eval_input_reader {
  label_map_path: "Data/CMT_Content/label_map.pbtxt"
  shuffle: false
  num_epochs: 1
  tf_record_input_reader {
    input_path: "Data/CMT_Content/Validation/cmt_content_validation.tfrecord"
  }
}

And I use exporter_main_v2.py to export the models

4. Expected behavior

Output a folder like normal Faster R-CNN without number_of_stages: 1 in the config file

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): Windows 10
• TensorFlow installed from (source or binary): Source
• TensorFlow version (use command below): 2.4
• Python version: 3.7
• CUDA/cuDNN version: 11.0/8.0.4
• GPU model and memory: GTX 1050