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

Posted

技术标签:

【中文标题】如何向预训练的对象检测模型添加其他类并训练它以检测所有类(预训练 + 新)?【英文标题】:How to add additional classes to a pre-trained object detection model and train it to detect all of the classes (pre-trained + new)? 【发布时间】:2020-01-15 21:16:18 【问题描述】:

我关注了这个博客 --> 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.

它仍然只检测新类。

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

【问题讨论】:

我似乎也找不到这个问题的明确答案。运气好吗? 不幸的是你不能。进行分类的密集层经过训练以输出特定类别。你不能指望网络输出它在训练期间没有看到的类。 【参考方案1】:

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

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

第一个选项只训练检测头,而不是网络的主干 - 这意味着主干权重在您的模型和原始模型之间分离。

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

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

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

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

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

这是一个关于如何从一个图形中提取权重并将它们组合到一个新图形中的教程(适用于 TF1)TensorFlow: saving/restoring and mixing multiple models

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

【讨论】:

以上是关于如何向预训练的对象检测模型添加其他类并训练它以检测所有类(预训练 + 新)?的主要内容,如果未能解决你的问题,请参考以下文章

使用重新训练的 Tensorflow 对象检测模型使用 snpe 进行 pb 到 dlc 转换失败

如何在 Tensorflow 对象检测 api 中评估预训练模型

如何训练具有可变输出大小的对象检测模型?

如何对经过训练的目标检测模型进行剪枝?

人工智能神经网络不训练检测模型,如何进行对象的检测与追踪?

我可以在已经分割的小图像上训练 YOLO 并在大图像上测试它以进行检测吗?