如何在预训练的对象检测模型中添加其他类并对其进行训练以检测所有类(预训练的+新的)?

时间:2019-09-14 06:31:42

标签: python tensorflow deep-learning object-detection

我关注了这个博客-> https://medium.com/@teyou21/training-your-object-detection-model-on-tensorflow-part-2-e9e12714bdf,并建立了一个SSD Mobilenet模型,该模型在名为“ ssd_mobilenet_v2_quantized_coco”的COCO数据集上进行了预训练。

这里发生的是,它可以完美地检测到我的新课程,但是我也想包括经过预训练的课程。

我尝试将类数更改为96(90个预训练+ 6个新类),并使用来自COCO数据集的所有标签的名称和相应ID编辑了“ labelmap.pbtxt”,并在该类中添加了新类最后从ID 91-96开始。

它仍然仅仅检测新类。

我该怎么做才能同时检测预训练课程和新课程?

1 个答案:

答案 0 :(得分:3)

这取决于您如何使用预先训练的重量:

  1. 用于迁移学习(我认为您发送的链接就是他们的工作)
  2. 使用有一个适合整个模型的起点。

第一个选项仅训练检测头,而不训练网络的主干-这意味着主干权重位于模型和原始模型之间。

在第二个选项中,您训练所有网络,主干+检测头-这意味着您有两种不同的模型

如果您使用的是第二个选项,那么执行所需操作的唯一方法是加载两个网络,并对原始网络运行一次镜像,然后对新网络运行一次镜像。然后,您可以合并结果。

如果使用第一个选项,则可以执行以下操作:

  1. 训练网络上的数据并保存新的检测头权重。
  2. 创建一个具有相同主干但有两个检测头的新网络: 一个带有原始砝码,另一个带有新砝码。

这个想法是因为因为两者的主干都是相同的,所以我们可以使用主干提取图像的特征,然后为每个检测头提供特征。

这是一本有关如何从一个图形中提取权重并将其合并为一个新的权重的教程(这是针对TF1的)TensorFlow: saving/restoring and mixing multiple models

在这里您可以了解如何保存和还原模型的一部分-save-and-restore-a-subset-of-variables