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

Question

我训练了一个 Faster R-CNN来自 TF Object Detection API并使用 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

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

1. 包含范围 FirstStageFeatureExtractor 中每个变量的模型，它接受图像(或序列化的 tf.data.Example)作为输入，并输出特征图和 RPN 提案。

2. 包含范围 SecondStageFeatureExtractor 和 SecondStageBoxPredictor 中的每个变量的模型，它接受特征图和 RPN 建议作为输入，并输出边界框预测和分数.

我基本上希望能够调用_predict_first_stage和 _predict_second_stage分别在我的输入数据上。

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

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

2020 年 6 月 7 日编辑:对于模型 1，我可以提取 detection_features，如 this question ，但我仍然不确定模型 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。