【Hugging Face】datasets 库：加载、处理和分享大规模数据集
Hugging Face Datasets 库
Hugging Face 的 datasets 库是一个 轻量级、高性能 的库，用于 加载、处理和分享大规模数据集，特别适用于 自然语言处理（NLP）、计算机视觉（CV）和语音任务。
1. 为什么使用 Datasets？
在深度学习中，处理大规模数据集通常面临以下挑战：
- 数据集太大，无法一次性加载到内存
- 不同任务的数据格式不统一
- 数据预处理和转换较慢
- 需要快速流式加载数据
datasets 库提供了解决方案：
- 支持数千个开源数据集，可直接加载 Hugging Face Hub 上的数据
- 流式加载（Streaming），可处理超大规模数据
- 兼容 PyTorch、TensorFlow、JAX
- 原生支持 Apache Arrow，加快数据加载
- 提供强大的数据预处理、切分、格式转换功能
2. 安装 Datasets
可以直接使用 pip 安装：
pip install datasets
3. 主要功能
3.1. 直接加载预训练数据集
from datasets import load_dataset
dataset = load_dataset("imdb")
print(dataset)
输出：
DatasetDict({
train: Dataset({
features: ['text', 'label'],
num_rows: 25000
})
test: Dataset({
features: ['text', 'label'],
num_rows: 25000
})
})
这表示 imdb 数据集有 train 和 test 两个子集，每个样本包含 text 和 label。
3.2. 访问数据
你可以像操作字典一样访问不同的数据集：
train_data = dataset["train"]
print(train_data[0]) # 访问第一条数据
输出：
{'text': 'This is a great movie!', 'label': 1}
如果要取出所有数据：
texts = dataset["train"]["text"]
labels = dataset["train"]["label"]
print(texts[:5]) # 查看前5条文本
3.3. 过滤和转换数据
如果你只想保留正面评论（label = 1），可以使用 .filter()：
positive_reviews = dataset["train"].filter(lambda x: x["label"] == 1)
print(positive_reviews)
如果你想对 text 进行转换，例如变成小写：
dataset = dataset.map(lambda x: {"text": x["text"].lower()})
这样 text 字段的所有内容都会变成小写。
3.4. 数据集切分
你可以手动切分训练和测试数据：
train_test_split = dataset["train"].train_test_split(test_size=0.2)
print(train_test_split)
这样 train_test_split 就会包含 train 和 test 两部分数据。
3.5. 数据批量加载
如果数据集很大，你可以使用 batch_size 进行批量加载：
dataset = dataset.map(preprocessing_function, batched=True)
这可以大幅加快数据处理速度。
3.6. 数据格式转换
datasets 允许数据集转换成 PyTorch 或 TensorFlow 兼容的格式：
dataset.set_format(type="torch", columns=["text", "label"])
或者：
dataset.set_format(type="tensorflow", columns=["text", "label"])
这样就可以直接用于 PyTorch 或 TensorFlow 训练。
3.7. 流式加载超大数据集
如果数据集太大，无法一次性加载到内存，可以使用 Streaming 模式：
dataset = load_dataset("c4", split="train", streaming=True)
for sample in dataset:
print(sample)
break
这样 datasets 会 逐行读取数据，不会占用大量内存。
4. 自定义数据集
如果你有自己的数据集（例如 CSV、JSON、TXT），可以使用 load_dataset 直接加载。
4.1. 加载 CSV 数据
dataset = load_dataset("csv", data_files="my_data.csv")
print(dataset)
4.2. 加载 JSON 数据
dataset = load_dataset("json", data_files="my_data.json")
4.3. 加载本地文本数据
dataset = load_dataset("text", data_files="my_text.txt")
5. 适用于不同任务的数据集
5.1. NLP 任务
可以加载 SQuAD 问答数据集：
dataset = load_dataset("squad")
5.2. 计算机视觉任务
可以加载 CIFAR-10 数据集：
dataset = load_dataset("cifar10")
5.3. 语音任务
可以加载 Common Voice 语音数据集：
dataset = load_dataset("mozilla-foundation/common_voice_11_0", "zh-CN")
6. Hugging Face Datasets 和 PyTorch Dataset 的对比
| 功能 | datasets |
torch.utils.data.Dataset |
|---|---|---|
| 直接加载预训练数据集 | 是 | 否 |
| 高效数据预处理 | 是 | 否 |
| 自动流式加载 | 是 | 否 |
| 兼容 PyTorch/TensorFlow | 是 | 否 |
| 轻松支持大规模数据 | 是 | 否 |
如果你使用 PyTorch 的 DatasetDataLoader ，但希望用 datasets 加速，可以这样转换：
from torch.utils.data import DataLoader
dataloader = DataLoader(dataset["train"], batch_size=16, shuffle=True)
7. 总结
datasets是 Hugging Face 提供的 高效数据加载和处理库，支持 NLP、CV、语音等任务。- 可以直接加载数千个 Hugging Face 预训练数据集，例如
imdb、squad、cifar10。 - 支持 CSV、JSON、TXT 等本地数据格式，可以轻松加载自定义数据集。
- 支持 PyTorch 和 TensorFlow，并提供流式加载，适用于超大规模数据。
- 提供强大的数据切分、过滤、转换功能，让数据处理变得简单高效。