如何向预训练的对象检测模型添加额外的类并训练它检测所有类(预训练 + 新的)?

Aad*_*dar 9 python object-detection deep-learning tensorflow

我关注了这个博客 --> 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”,最后从 ids 添加新类91 - 96。

它仍然只检测新类。

我应该怎么做才能检测到预训练的和新的类?

Ami*_*ani 5

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

  1. 用于迁移学习(我认为在您发送的链接中这就是他们所做的)
  2. 使用有一个拟合整个模型的起点。

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

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

如果在您的情况下您使用第二个选项,那么执行您想要的操作的唯一方法是加载两个网络并使用原始网络对图像运行推理一次,然后使用新网络运行推理。然后你结合你的结果。

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

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

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

这是一篇关于如何从一个图中提取权重并将其组合到一个新图中的教程(适用于 TF1)TensorFlow:保存/恢复和混合多个模型

在这里您可以阅读如何保存和恢复模型的一部分 - save-and-restore-a-subset-of-variables