重要提示

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

数据准备以使用 Megatron-Energon 数据加载器#

Megatron-Energon 数据加载器旨在高效处理大规模分布式训练环境，特别是对于多模态模型。本文档概述了如何准备所需格式的数据以与 NeMo 框架一起使用。

有关 Megatron-Energon 的详细信息，请参阅 Megatron-Energon 简介。有关 WebDataset 格式的详细信息，请参阅 WebDataset 文档。

使用 Megatron-Energon 数据加载器的第一步是以支持的格式准备数据。这包括将原始数据集转换为 WebDataset (WDS) 格式，并生成 Megatron-Energon 所需的元数据。

在此示例中，我们将展示如何将 LLaVA-Pretrain 数据集（JSON 格式）转换为 WebDataset 格式，然后准备它以供 NeMo 中的 Energon 数据模块使用。

步骤 1：下载数据集#

使用以下链接从 Hugging Face 下载数据集

以下是 JSON 结构的示例

[
  {
    "id": "GCC_train_002582585",
    "image": "GCC_train_002582585.jpg",
    "conversations": [
      {
        "from": "human",
        "value": "Provide a brief description of the given image.\n<image>"
      },
      {
        "from": "gpt",
        "value": "Olive oil is a healthy ingredient used liberally."
      }
    ]
  }
]

步骤 2：将 JSON 数据转换为 WebDataset 格式#

可以使用 Python 脚本将 JSON 数据集转换为 WebDataset 格式。您还可以参考 Megatron-LM 中提供的此示例：convert_llava_pretrain_to_wds.py。以下是将执行转换的 Python 脚本示例。根据数据的结构和键修改脚本。该脚本可以从任何目录运行。只需修改输入路径（‘llava_pretrain_dir’）和下面的输出路径。转换后的数据存储在输出路径 <path_to_LLaVA-Pretrain>/wds 中。

import json
import os
import webdataset as wds
from tqdm import tqdm

# Set the path to the LLaVA-Pretrain dataset directory
llava_pretrain_dir = '<path_to_LLaVA-Pretrain>'

# Paths to the dataset files
json_file = os.path.join(llava_pretrain_dir, 'blip_laion_cc_sbu_558k.json')
output_path = os.path.join(llava_pretrain_dir, 'wds')

if not os.path.exists(output_path):
    os.mkdir(output_path)

# Load data
with open(json_file, 'r') as f:
    data = json.load(f)

# Convert JSON to WebDataset
with wds.ShardWriter(os.path.join(output, 'pretrain-%d.tar'), maxcount=10000) as shard_writer:
    for entry in tqdm(data):
        with open(os.path.join(llava_pretrain_dir, entry['image']), "rb") as img_file:
            image_data = img_file.read()
        sample = {
            "__key__": entry['id'],
            "jpg": image_data,
            "json": json.dumps(entry['conversations']).encode("utf-8"),
        }
        shard_writer.write(sample)

print(f"Dataset successfully converted to WebDataset format.")

步骤 3：为 Megatron-Energon 生成元数据#

将数据集转换为 WebDataset 格式后，必须为 Megatron-Energon 数据加载器生成元数据。在包含 WebDataset 文件的目录中运行以下命令：导航到上述输出目录路径以运行以下命令。

energon prepare .

或

energon prepare <path_to_LLaVA-Pretrain>/wds

此命令创建一个 .nv-meta 文件夹，其中包含 Megatron-Energon 数据加载器所需的必要元数据。

有关更多详细信息，请参阅 Megatron-Energon 文档中的数据准备部分。

数据集转换和元数据生成后，您的文件夹将如下所示

.
├── .nv-meta
│   ├── dataset.yaml
│   ├── split.yaml
├── pretrain-0.tar
├── pretrain-16.tar
├── pretrain-16.tar.idx
├── pretrain-23.tar
├── pretrain-23.tar.idx
├── pretrain-31.tar
├── pretrain-31.tar.idx
├── pretrain-39.tar
├── pretrain-39.tar.idx
├── pretrain-47.tar
├── pretrain-47.tar.idx
├── pretrain-55.tar
└── pretrain-55.tar.idx

.nv-meta 目录内容#

.nv-meta 文件夹包含 Megatron-Energon 数据加载器必不可少的元数据文件

.nv-meta/
├── dataset.yaml
└── split.yaml

dataset.yaml 的示例内容#

dataset.yaml 文件提供有关数据集结构和字段映射的信息。以下是一个示例

__class__: VQAWebdataset
__module__: megatron.energon
field_map:
  answers: json[1][value]
  context: json[0][value]
  image: jpg

其他参考资料#

有关数据集格式和示例的更多信息，请参阅：Megatron-Energon 数据集格式

有关创建自定义样本和自定义数据加载器的详细信息，请参阅：高级数据格式文档。

在本页中