【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. 总结

1. datasets 是 Hugging Face 提供的 高效数据加载和处理库，支持 NLP、CV、语音等任务。
2. 可以直接加载数千个 Hugging Face 预训练数据集，例如 imdb、squad、cifar10。
3. 支持 CSV、JSON、TXT 等本地数据格式，可以轻松加载自定义数据集。
4. 支持 PyTorch 和 TensorFlow，并提供流式加载，适用于超大规模数据。
5. 提供强大的数据切分、过滤、转换功能，让数据处理变得简单高效。