datasets开发者指南：构建自定义数据集的最佳实践-优快云博客

🤗 datasets开发者指南：构建自定义数据集的最佳实践

【免费下载链接】datasets 🤗 The largest hub of ready-to-use datasets for ML models with fast, easy-to-use and efficient data manipulation tools 项目地址: https://gitcode.com/gh_mirrors/da/datasets

你是否还在为数据格式不兼容、内存溢出、处理速度慢而烦恼？本文将带你掌握构建自定义数据集的完整流程，从文件组织到高效加载，从元数据管理到格式转换，让你的机器学习项目数据处理环节效率提升10倍。读完本文，你将能够：

使用Folder-Based Builders快速创建图像/音频数据集
掌握从Python字典和生成器构建数据集的技巧
优化大数据集的内存使用和加载速度
正确处理元数据和标签系统
无缝集成PyTorch/TensorFlow等框架

为什么需要自定义数据集？

在机器学习项目中，数据准备往往占据整个开发周期的60%以上时间。🤗 datasets库提供了统一的数据接口和高效的数据处理工具，让你能够专注于模型开发而非数据处理。无论是处理图像、音频还是文本数据，自定义数据集都能帮助你：

统一数据格式，减少格式转换开销
利用内存映射（Memory Mapping）技术处理超大数据集
支持流式加载，降低内存占用
与主流机器学习框架无缝集成

官方文档：README.md 中详细介绍了datasets库的核心优势和基本使用方法。

环境准备

首先确保安装了datasets库及其相关依赖：

# 基础安装
pip install datasets

# 如需处理图像数据
pip install datasets[vision]

# 如需处理音频数据
pip install datasets[audio]

如果你使用PyTorch或TensorFlow，还需要安装对应的框架：

# PyTorch用户
pip install torch

# TensorFlow用户
pip install tensorflow

详细安装指南：docs/source/installation.md

快速开始：Folder-Based Builders

对于图像和音频数据集，datasets提供了开箱即用的Folder-Based Builders，只需简单的目录结构即可自动生成数据集。目前支持两种构建器：ImageFolder和AudioFolder，它们能够自动推断标签、生成特征和分割数据集。

图像数据集构建

假设你的图像数据按照以下结构组织：

pokemon/train/grass/bulbasaur.png
pokemon/train/fire/charmander.png
pokemon/train/water/squirtle.png

pokemon/test/grass/ivysaur.png
pokemon/test/fire/charmeleon.png
pokemon/test/water/wartortle.png

使用ImageFolder构建器只需一行代码：

from datasets import load_dataset

dataset = load_dataset("imagefolder", data_dir="/path/to/pokemon")

ImageFolder会自动将目录名作为标签，并根据目录结构生成训练/测试分割。支持的图像格式包括jpg、png等常见格式，完整列表可查看源码：src/datasets/packaged_modules/imagefolder/imagefolder.py#L39

音频数据集构建

音频数据集的构建方式类似，只需将imagefolder替换为audiofolder：

from datasets import load_dataset

dataset = load_dataset("audiofolder", data_dir="/path/to/audio_data")

AudioFolder支持wav、mp3、mp4等多种音频格式，解码通过ffmpeg实现，支持的格式列表：ffmpeg.org/ffmpeg-formats.html

目录结构解析

Folder-Based Builders的核心原理是通过目录结构推断数据集的结构和标签。以下是一个典型的目录结构示例：

dataset_root/
├── split1/          # 数据集分割（如train/validation/test）
│   ├── class1/      # 类别标签1
│   │   ├── file1.jpg
│   │   └── file2.jpg
│   └── class2/      # 类别标签2
│       ├── file3.jpg
│       └── file4.jpg
└── split2/
    ├── class1/
    └── class2/

这种结构会被自动解析为包含splits和labels的数据集。内部实现可参考：src/datasets/packaged_modules/folder_based_builder/folder_based_builder.py

从Python数据结构构建数据集

当你需要更灵活的数据处理时，可以从Python字典或生成器构建数据集。这两种方法各有优势，适用于不同场景。

从字典构建

对于中小型数据集，Dataset.from_dict()是最简单直接的方法：

from datasets import Dataset

data_dict = {
    "pokemon": ["bulbasaur", "squirtle", "charmander"],
    "type": ["grass", "water", "fire"],
    "hp": [45, 44, 39]
}

dataset = Dataset.from_dict(data_dict)
print(dataset[0])
# 输出: {"pokemon": "bulbasaur", "type": "grass", "hp": 45}

这种方法会将字典中的每个键作为特征列，值作为该列的数据。适合数据已经全部加载到内存中的情况。

从生成器构建

对于大型数据集，推荐使用Dataset.from_generator()，它通过迭代方式生成数据，避免一次性加载全部数据到内存：

from datasets import Dataset

def data_generator():
    for i in range(1000000):
        yield {
            "id": i,
            "value": i * 2,
            "text": f"Sample {i} with value {i*2}"
        }

# 创建常规数据集
dataset = Dataset.from_generator(data_generator)

# 或创建可迭代数据集（适合流式处理）
from datasets import IterableDataset
iterable_dataset = IterableDataset.from_generator(data_generator)

# 迭代访问
for sample in iterable_dataset.take(3):
    print(sample)

生成器方法特别适合处理无法全部放入内存的大型数据集，因为它会在磁盘上逐步生成数据并进行内存映射。详细实现见：src/datasets/arrow_dataset.py

高级特性：元数据与特征处理

元数据管理

当你的数据集需要额外的元数据（如 captions、transcriptions等）时，可以通过metadata.csv文件添加。元数据文件需要包含file_name列，用于关联主数据文件：

file_name,text,caption
bulbasaur.png,"There is a plant seed on its back","Grass type pokemon"
charmander.png,"It has a preference for hot things","Fire type pokemon"
squirtle.png,"When it retracts its long neck into its shell, it squirts out water","Water type pokemon"

加载带元数据的数据集：

dataset = load_dataset("imagefolder", data_dir="/path/to/pokemon")

系统会自动检测并加载元数据文件，支持的格式包括CSV、JSONL和Parquet。元数据处理逻辑在：src/datasets/packaged_modules/folder_based_builder/folder_based_builder.py#L368-L393

特征类型转换

datasets支持多种特征类型，包括Image、Audio、ClassLabel等，可以使用cast_column方法转换特征类型：

# 转换为图像特征
from datasets import Image
dataset = dataset.cast_column("image", Image())

# 转换为音频特征并指定采样率
from datasets import Audio
dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))

# 转换为分类标签
from datasets import ClassLabel
dataset = dataset.cast_column("label", ClassLabel(names=["grass", "fire", "water"]))

特征系统的核心实现位于：src/datasets/features/features.py

性能优化：处理大型数据集

内存优化

对于大型数据集，使用流式加载和迭代处理可以显著降低内存占用：

# 流式加载数据集
dataset = load_dataset("csv", data_files="large_file.csv", streaming=True)

# 迭代处理
for batch in dataset.iter(batch_size=1000):
    process_batch(batch)

流式处理不会将整个数据集加载到内存，而是按需加载数据块。详细说明见：docs/source/stream.mdx

数据格式转换

将数据集转换为Arrow格式可以提高加载速度和降低内存占用：

# 保存为Arrow格式
dataset.save_to_disk("my_dataset")

# 后续加载
from datasets import load_from_disk
dataset = load_from_disk("my_dataset")

Arrow格式支持内存映射，允许你像访问本地数组一样访问磁盘上的大型数据集。Arrow相关实现：src/datasets/arrow_reader.py

框架集成：PyTorch/TensorFlow

PyTorch集成

将数据集转换为PyTorch张量并创建DataLoader：

# 设置格式为PyTorch张量
dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])

# 创建DataLoader
from torch.utils.data import DataLoader
dataloader = DataLoader(dataset, batch_size=32)

# 训练循环
for batch in dataloader:
    inputs = {k: v.to(device) for k, v in batch.items()}
    outputs = model(**inputs)
    loss = outputs.loss
    loss.backward()

TensorFlow集成

类似地，可以将数据集转换为TensorFlow格式：

# 转换为TensorFlow数据集
tf_dataset = dataset.to_tf_dataset(
    columns=["input_ids", "attention_mask"],
    label_cols=["labels"],
    batch_size=32,
    shuffle=True
)

# 训练
model.fit(tf_dataset, epochs=3)

框架集成的详细文档：docs/source/use_with_pytorch.mdx 和 docs/source/use_with_tensorflow.mdx

实战案例：构建图像分类数据集

让我们通过一个完整案例，构建一个用于图像分类的自定义数据集：

pokemon_dataset/
├── train/
│   ├── grass/
│   │   ├── bulbasaur.png
│   │   └── ivysaur.png
│   ├── fire/
│   │   ├── charmander.png
│   │   └── charmeleon.png
│   └── water/
│       ├── squirtle.png
│       └── wartortle.png
├── validation/
│   ├── grass/
│   ├── fire/
│   └── water/
└── metadata.csv

加载数据集

from datasets import load_dataset

dataset = load_dataset("imagefolder", data_dir="pokemon_dataset")
print(dataset)
# 输出: 
# DatasetDict({
#     train: Dataset({
#         features: ['image', 'label'],
#         num_rows: 6
#     })
#     validation: Dataset({
#         features: ['image', 'label'],
#         num_rows: 6
#     })
# })

数据预处理

# 添加数据增强
from torchvision.transforms import Compose, ColorJitter, ToTensor

transform = Compose([
    ColorJitter(brightness=0.5, hue=0.5),
    ToTensor()
])

def preprocess(examples):
    examples["pixel_values"] = [transform(image.convert("RGB")) for image in examples["image"]]
    return examples

# 设置转换（实时应用，不修改原始数据）
dataset = dataset.with_transform(preprocess)

准备训练

# 转换为PyTorch格式
dataset.set_format("torch", columns=["pixel_values", "label"])

# 创建DataLoader
from torch.utils.data import DataLoader
train_loader = DataLoader(dataset["train"], batch_size=4, shuffle=True)
val_loader = DataLoader(dataset["validation"], batch_size=4)

# 现在可以直接用于模型训练了

总结与最佳实践

1.** 数据组织 **- 对于图像/音频数据，优先使用Folder-Based结构，简化标签管理

使用清晰的目录结构区分数据集分割（train/val/test）
元数据使用单独文件管理，保持主数据干净

2.** 性能优化 **- 大型数据集使用生成器或流式加载

频繁访问的数据集转换为Arrow格式
使用with_transform进行实时数据增强，避免数据复制

3.** 兼容性 **- 转换特征类型以匹配模型输入要求

使用框架专用格式转换（set_format）确保兼容性
对于分布式训练，考虑使用Dataset.shard()拆分数据

4.** 可维护性 **- 为自定义数据集编写详细文档

使用标准特征命名（如"image"、"audio"、"text"）
版本控制数据集，记录变更历史

通过本文介绍的方法，你可以构建高效、灵活且易于维护的自定义数据集，为你的机器学习项目打下坚实基础。更多高级用法和最佳实践，请参考：docs/source/how_to.md

下一步学习资源

官方教程：docs/source/tutorial.md
API参考：docs/source/package_reference/main_classes.mdx
社区案例：notebooks/README.md
问题排查：docs/source/troubleshoot.mdx

掌握自定义数据集构建技能，让你的机器学习项目数据处理环节效率倍增，专注于模型创新而非数据琐事！

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考