3、使用小数据集快速原型制作 | NVIDIA 开发者
利用迁移学习调整和助力 AI 工作流的各种方法
NVIDIA TAO 工具套件和预训练的模型如何改变您的开发工作
3. 使用小型数据集快速进行原型设计
3.1 为 AI 模型训练采集数据所面临的挑战
在 AI 领域，“输入的是垃圾，输出的也必然是垃圾”这句老话非常适用。如果从零开始训练模型，您需要有大型的高质量数据集。下面列出了数据采集存在的一些挑战：
- 采集和标记数据不仅费时费力，而且成本高昂
- 使用小型数据集可能会导致模型性能很差，因为它们缺少成功训练稳健的模型所需的多样性
- 在某些情况下，数据可能根本无法获得或受限（例如，患者医疗 X 光片和医学扫描件）
举例来说，检测 PCB 装配线上的缺陷是一项关键任务，有助于确保质量，并发现装配过程中的错误。但是，PCB 装配中出现缺陷的比率非常低，可能需要数月甚至数年的时间才能采集到足够多的图像来训练出准确的模型。在尝试检测异常情况时，许多行业都面临着需要采集足够多的数据这一挑战。
下面是克服这些挑战的一些方法：
- 生成合成数据：合成数据是计算机仿真或算法生成的带标注的信息。
- 数据增强：增强数据集可提高数据的多样性和随机性，从而实现模型泛化，而这有助于模型更准确地处理之前未见过的数据。
与采集更多数据相比，这两种方法都可以显著降低成本并提升速度。在本实验中，我们将了解 TAO 工具套件中的数据增强功能。
3.2 什么是数据增强？
数据增强是指：采用现有数据集，并在空间域和颜色域中应用转换，以创建与原始图像相似但又足够不同的新图像，从而泛化模型并提高多样性。为了确定更为有效的增强技术类型，我们已经进行了大量研究。常见的转换包括转化、旋转和颜色偏移。
如果使用小型数据集对模型进行训练，模型会开始记忆数据中的模式，而不是学习解决问题所需的特征。通过应用数据增强来增大数据集的规模可以提高数据的复杂性，并迫使模型实现泛化而非进行记忆。这有助于减少训练集的过拟合，并提高在处理之前未见过的图像时的性能。模型可能会遇到处于不同光照条件、位置和方向的对象时，增强功能尤为有用。
对数据集应用增强功能既可以离线进行，也可以在线进行。离线增强在训练之前应用，并且会在存储空间中创建应用了转换的新图像。这样便可以使用更多单一图像数据，通常能够使模型在更少迭代次数内实现收敛。在线增强功能会对在训练中使用的每张图像动态应用随机转换。这意味着，无需存储额外的图像，也无需额外的磁盘空间。此外，由于每次应用转换后都会创建一张唯一的图像，因此这还可以使模型不断使用新图像进行训练。使用在线增强功能训练模型时，由于模型会不断“看到”新图像，因此可能需要更多次迭代才能实现收敛。
3.3 通过 TAO 工具套件应用数据增强
TAO 工具套件既支持在线数据增强，也支持离线数据增强。您可以通过以下方式执行离线增强：通过一个配置文件，然后使用命令行界面来生成图像。通过该配置，您可以自定义空间、颜色和模糊增强。您还可以自定义在线增强，以便指定训练模型时要应用的空间和颜色增强范围。
若要在 TAO 工具套件中使用增强功能，建议的方法是，首先应用离线增强来增大数据集的规模，然后将训练配置为使用在线增强，以进一步提高数据集的复杂性。如增强任务结果中所示，在结合使用这两种类型的增强功能后，模型能够“看到”大量不同的图像，并实现更好的性能。
下表显示了 TAO 工具套件支持的不同增强技术。
| 空间 | 离线 | 在线 | 颜色 | 离线 | 在线 | |
|---|---|---|---|---|---|---|
| 旋转 | 色调旋转 | |||||
| 翻转 | 饱和度偏移 | |||||
| 转化 | 对比度 | |||||
| 剪切 | 亮度 | |||||
| 缩放 | 颜色偏移 | |||||
| 模糊 | ||||||
| 表 2. 显示了 TAO 中所有可能的空间和颜色增强 | ||||||
3.4 结果
为了展示数据增强的优势，我们以使用仅包含 100 张图像的 PCB 缺陷数据集进行训练为例。该模型的任务是从 PCB 图像中检测六种类型的缺陷。这项任务的关键性能指标是平均精度均值 (mAP)，该指标可衡量模型在缺陷周围放置的边界框相较于真值的准确度。
在这项任务中，我们仅使用 100 张未增强的图像对模型进行了训练，然后应用离线增强来创建数据集，其规模是原始数据集的 10 倍和 20 倍。接下来，我们分别在应用在线增强和不应用在线增强的情况下使用这些离线增强的数据集进行训练，以便比较。
如果在不应用任何增强的情况下使用 100 张图像进行训练，模型在处理测试集时仅能达到 36% 的 mAP。但是，在应用这两种形式的增强后，模型有了显著改善，能够达到将近 79% 的 mAP。这项任务表明，在数据集小到只有 100 张图像时，增强功能在提高模型质量方面非常有效。
初始原型模型只使用了 100 张图像。然后，我们添加了更多数据以实现更高的准确度。接下来，我们使用 500 张带标注的图像执行了相同的任务。在没有应用任何增强的情况下，数据增加使 mAP 达到了 82%。此外，在应用上述两种增强的情况下，mAP 提升到了 95% 以上。
您可以将 TAO 工具套件提供的增强功能应用于任何规模的数据集。如果想要使用小型数据集快速对机器学习应用进行原型设计，或想要进一步改进高度准确的模型，该工具是理想之选。
这项任务使用了 PCB 缺陷数据集 2。TAO 任务 GitHub 资料库中提供了重现这些结果所需的所有代码和步骤。
1 Shorten, C., Khoshgoftaar, T.M. A survey on Image Data Augmentation for Deep Learning. J Big Data 6, 60 (2019). https://journalofbigdata.springeropen.com/articles/10.1186/s40537-019-0197-0
2 R. Ding, L. Dai, G. Li and H. Liu, "TDD-net: a tiny defect detection network for printed circuit boards," in CAAI Transactions on Intelligence Technology, vol. 4, no. 2, pp. 110-116, 6 2019, doi: 10.1049/trit.2019.0019.