重要提示

您正在查看 NeMo 2.0 文档。此版本对 API 和新库 NeMo Run 进行了重大更改。我们目前正在将 NeMo 1.0 的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

图像-文本对数据集#

图像-文本对数据集通常用于训练生成式文本到图像模型或 CLIP 模型。NeMo Curator 支持基于 WebDataset 文件格式读取和写入数据集。此格式允许 NeMo Curator 使用元数据（包括嵌入和分类器分数）注释数据集。其分片格式也使其更容易将工作分发给处理数据集的不同工作进程。

文件格式#

以下是以 WebDataset 格式的数据集目录的示例。

dataset/
├── 00000.tar
│   ├── 000000000.jpg
│   ├── 000000000.json
│   ├── 000000000.txt
│   ├── 000000001.jpg
│   ├── 000000001.json
│   ├── 000000001.txt
│   └── ...
├── 00001.tar
│   ├── 000010000.jpg
│   ├── 000010000.json
│   ├── 000010000.txt
│   ├── 000010001.jpg
│   ├── 000010001.json
│   ├── 000010001.txt
│   └── ...
├── 00002.tar
│   └── ...
├── 00000.parquet
├── 00001.parquet
└── 00002.parquet

确切的格式假定一个包含分片 .tar、.parquet 和（可选）.idx 文件的单个目录。每个 tar 文件应具有唯一的整数 ID 作为其名称（00000.tar、00001.tar、00002.tar 等）。tar 文件应包含 .jpg 文件中的图像、.txt 文件中的文本标题和 .json 文件中的元数据。数据集的每个记录都由唯一的 ID 标识，该 ID 是分片 ID 以及记录在分片内的偏移量的混合。例如，第 43 个分片的第 32 个记录将在 00042.tar 中，并具有图像 000420031.jpg、标题 000420031.txt 和元数据 000420031.json（假设从零开始索引）。

除了 tar 文件集合外，NeMo Curator 的 ImageTextPairDataset 期望根目录中存在 .parquet 文件，这些文件遵循与分片相同的命名约定（00042.tar -> 00042.parquet）。每个 Parquet 文件应包含每个记录的元数据的聚合表格形式，Parquet 文件中的每一行对应于该分片中的记录。Parquet 文件和 JSON 文件中的元数据都必须包含一个唯一的 ID 列，该列与其记录 ID 相同（在我们的示例中为 000420031）。

读取#

可以使用 ImageTextPairDataset.from_webdataset() 读取数据集

  • path="/path/to/dataset" 应指向 WebDataset 的根目录。

  • id_col="key" 让我们知道数据集中的唯一 ID 列名为 “key”。

更完整的参数列表可以在 API 参考 中找到。

写入#

有两种方法可以写入图像数据集。第一种方法仅保存元数据，而第二种方法将重新分片 tar 文件。两者都会触发您之前设置为运行的所有任务的计算。

save_metadata() 将仅将分片 Parquet 文件保存到目标目录。它不会修改 tar 文件。有两个可选参数

  • path 允许您更改数据集的保存位置。默认情况下，它将覆盖原始 Parquet 文件。

  • columns 允许您仅保存元数据的子集。默认情况下，将保存所有元数据。

to_webdataset() 将重新分片 WebDataset，使其仅包含在 filter_column 中具有 True 值的元素。重新分片可能需要一段时间，因此通常应仅在整理管道结束时且您准备好导出数据集以进行训练时执行此操作。

更完整的参数列表可以在 API 参考 中找到。

索引文件#

NeMo Curator 使用 DALI 从 tar 文件加载图像数据。为了加快数据加载速度，您可以在数据集中提供 .idx 文件。索引文件必须由 DALI 的 wds2idx 工具生成。有关更多信息，请参阅 DALI 文档。每个索引文件必须遵循与 tar 文件相同的命名约定 (00042.tar -> 00042.idx)。

在本页中