如何在 TF Object Detection 2.0 中分别加载已保存的 Faster R-CNN 的两个阶段？

Question

我训练了一个Faster R-CNN from the TF Object Detection API and saved it using export_inference_graph.py。我有以下目录结构：

weights
|-checkpoint
|-frozen_inference_graph.pb
|-model.ckpt-data-00000-of-00001
|-model.ckpt.index
|-model.ckpt.meta
|-pipeline.config
|-saved_model
|--saved_model.pb
|--variables

我想分别加载模型的第一阶段和第二阶段。也就是说，我想要以下两个型号：

包含范围 FirstStageFeatureExtractor 中每个变量的模型，它接受图像（或序列化 tf.data.Example）作为输入，并输出特征图和 RPN 提议。
包含范围 SecondStageFeatureExtractor 和 SecondStageBoxPredictor 中每个变量的模型，它接受特征图和 RPN 建议作为输入，并输出边界框预测和分数。

我基本上希望能够对我的输入数据单独调用 _predict_first_stage and _predict_second_stage。

目前我只知道如何加载整个模型：

model = tf.saved_model.load("weights/saved_model")
model = model.signatures["serving_default"]

编辑 2020 年 6 月 7 日：对于模型 1，我可能能够像 this question 那样提取 detection_features，但我仍然不确定模型 2。

Answer 1

这在对象检测仅与 TF1 兼容时更加困难，但现在在 TF2 中非常简单。 this colab.

中有一个很好的例子

from object_detection.builders import model_builder
from object_detection.utils import config_util

# Set path names
model_name = 'centernet_hg104_512x512_kpts_coco17_tpu-32'
pipeline_config = os.path.join('models/research/object_detection/configs/tf2/',
                                model_name + '.config')
model_dir = 'models/research/object_detection/test_data/checkpoint/'

# Load pipeline config and build a detection model
configs = config_util.get_configs_from_pipeline_file(pipeline_config)
model_config = configs['model']
detection_model = model_builder.build(model_config=model_config, 
                                     is_training=False)

# Restore checkpoint
ckpt = tf.compat.v2.train.Checkpoint(
      model=detection_model)
ckpt.restore(os.path.join(model_dir, 'ckpt-0')).expect_partial()

从这里可以调用 detection_model.predict() 和相关方法，例如 _predict_first_stage 和 _predict_second_stage。

如何在 TF Object Detection 2.0 中分别加载已保存的 Faster R-CNN 的两个阶段？

How do I load the two stages of a saved Faster R-CNN separately in TF Object Detection 2.0?

object-detection

tensorflow

object-detection-api

tensorflow2.0