在自定义数据集上训练 Faster-RCNN 模型时加载检查点

Posted 2023-03-13

【中文标题】在自定义数据集上训练 Faster-RCNN 模型时加载检查点

【英文标题】：Loading checkpoints while training a Faster-RCNN model on a custom dataset

【发布时间】：2021-12-10 02:16:09

【问题描述】：

我正在尝试使用 Faster-RCNN 架构（准确地说是Faster R-CNN ResNet50 V1 640x640，来自here）加载检查点并填充模型权重。我正在尝试为该网络加载权重，类似于它在example notebook for RetinaNet，他们在其中执行以下操作：

fake_box_predictor = tf.compat.v2.train.Checkpoint(
    _base_tower_layers_for_heads=detection_model._box_predictor._base_tower_layers_for_heads,
    _box_prediction_head=detection_model._box_predictor._box_prediction_head,
)

fake_model = tf.compat.v2.train.Checkpoint(
          _feature_extractor=detection_model._feature_extractor,
          _box_predictor=fake_box_predictor
)

ckpt = tf.compat.v2.train.Checkpoint(model=fake_model)
ckpt.restore(checkpoint_path).expect_partial()

我正在尝试为我想要使用的 Faster-RCNN 网络提供类似的检查点加载机制，但 _base_tower_layers_for_heads、_box_prediction_head 等属性仅适用于示例中使用的架构，而不适用于还有什么。

对于我的特定用例，我也找不到关于使用 Checkpoint 填充模型的哪些部分的文档。非常感谢任何有关如何解决此问题的帮助！

【参考方案1】：

正如您所说，您遇到的主要问题是您没有要在其上进行迁移学习的层的层张量。 这是 Zoo 中 Faster R-CNN ResNet50 V1 640x640 副本的原始实现的一部分。他们没有命名图层，或者他们确实命名了但没有公布名称（或代码）。 要解决这个问题，您需要找出要保留哪些层以及要重新学习哪些层。您可以使用 (ref) 打印出网络中的所有层：

[n.name for n in tf.get_default_graph().as_graph_def().node]

层的名称可以手动添加，但 tf 为每个节点保留默认名称。这个列表可能很长而且很累，但是你需要找到张量来开始你的迁移学习。因此，您需要按照列表并尝试了解要冻结的层以及要继续学习过程的层。冻结图层 (ref)：

if layer.name == 'layer_name':
    layer.trainable = False