4.1 为现有模型添加新分类所面临的挑战

如果您从可公开获得的模型开始操作，其中的分类可能不同于您的应用所需的分类。若要自定义模型以添加或移除分类，您需要知道并了解模型架构。这意味着要查看成千上万行代码以进行所需的更改，并且会涉及更改模型模板、数据加载器和损失函数。迁移学习任务之一是为现有模型添加新分类或从现有模型中移除一些分类。

比如，为现有的人员检测模型添加头盔分类。任务是检测人员和头盔。在训练这个新模型以添加头盔分类时，您必须提供同时包含人员和头盔且已正确标注的数据集。如果只提供头盔数据而不提供人员数据，模型在检测头盔方面表现良好，但在检测人员方面表现欠佳。一般来说，在迁移学习中，添加新分类时，您需要具有代表性的数据，并且这些数据要涵盖现有分类以及新分类。

4.2 在 TAO 工具套件中添加新分类

预训练的 NVIDIA 模型可帮助推理并生成预训练时所用的常见对象的标签。这样，您只需为自定义分类添加标签。然后，可以使用整个数据集对模型进行训练，在保持原始功能的同时，能够处理自定义分类。

以前面提到的人员和头盔为例，如果我们只拥有带标签的头盔数据集，可以使用 NVIDIA PeopleNet 模型中的推理功能来标注人员和面部。NVIDIA PeopleNet 模型已使用数以百万计的面部和人员图像进行训练，因此无需从零开始学习这些分类。我们首先使用 PeopleNet 进行推理，创建人员和面部标签，然后将推理出的标签与现有的头盔分类标签整合在一起。

注意：如果您使用 PeopleNet 生成人员和面部分类的真值，请注意这些分类的误报或漏报情况。可能需要进行一些手动清理。

鉴于 PeopleNet 模型已经能够高度准确地检测人员和面部，您只需训练头盔分类即可。为此，在训练规格中，我们要为头盔分类添加一个较高的分类权重，为人员和面部添加一个较低的分类权重。这有助于针对新分类更准确地训练模型，同时还能保持模型的人员和面部检测能力。

4.3 结果

在这项任务中，我们使用了开源头盔检测数据集，其中有 611 张图像用于训练，152 张图像用于验证。训练配置如下：为新的头盔分类设置 0.8 的权重，为人员和面部分类设置 0.1 的权重。这是为了帮助模型专注于学习新的头盔分类。然后，我们使用人员、头盔和面部数据集对 PeopleNet 模型进行了训练，并且在 100 次迭代内，模型在头盔分类方面便达到了 80% 的 AP。

采用预训练的模型，借助该模型进行数据集推理，然后再对其进行重新训练以包含自定义分类和原始分类，这种总体思路并不仅局限于 PeopleNet。在 TAO 中，您可以将该技术与任何预训练的模型结合使用。

这项任务使用了 Kaggle 头盔检测数据集。您可以在 TAO 任务 GitHub 资料库中找到完整的任务实现以及分步指南。

³ https://www.kaggle.com/andrewmvd/helmet-detection