hf datasets库

Hugging Face Datasets 库核心教程笔记

一、核心定位与学习前提

1. 库的目标：提供简单高效的方式加载、预处理、分享各类数据集，支持文本、图像、音频等多模态数据，兼容 PyTorch、TensorFlow 等主流机器学习框架。
2. 学习前提：具备基础 Python 知识，了解至少一种机器学习框架（如 PyTorch/TensorFlow）。

二、从 Hub 加载数据集

1. 先验调研：加载数据集构建器（无需下载数据）

使用 load_dataset_builder() 查看数据集元信息，提前评估是否符合需求：

from datasets import load_dataset_builder
ds_builder = load_dataset_builder("cornell-movie-review-data/rotten_tomatoes")

# 查看数据集描述
print(ds_builder.info.description)
# 查看数据集特征（字段名与类型）
print(ds_builder.info.features)

2. 加载数据集核心方法

使用 load_dataset() 下载并生成数据集，核心参数为数据集名称、配置、拆分：

from datasets import load_dataset
# 加载指定拆分
dataset = load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
# 不指定拆分，返回 DatasetDict（包含所有拆分）
dataset_dict = load_dataset("cornell-movie-review-data/rotten_tomatoes")

3. 关键概念：拆分（Splits）

• 拆分是数据集的子集，常见类型为 train/validation/test。
• 使用 get_dataset_split_names() 查看数据集支持的拆分：

  from datasets import get_dataset_split_names
  splits = get_dataset_split_names("cornell-movie-review-data/rotten_tomatoes")
  

• 加载指定拆分返回 Dataset 对象，不指定则返回 DatasetDict（键为拆分名称，值为对应 Dataset）。

4. 关键概念：配置（Configurations）

• 部分数据集包含多个子数据集（如多语言音频数据集），子数据集称为配置。
• 使用 get_dataset_config_names() 查看所有可用配置：

  from datasets import get_dataset_config_names
  configs = get_dataset_config_names("PolyAI/minds14")
  

• 加载时需显式指定配置，否则会抛出 ValueError：

  mindsFR = load_dataset("PolyAI/minds14", "fr-FR", split="train")
  

三、数据集对象操作（Dataset 与 IterableDataset）

1. 两类数据集对象对比

特性	Dataset	IterableDataset
加载方式	常规加载（全量加载/内存映射）	streaming=True 参数启用流式加载
访问方式	支持随机访问（索引、切片）	仅支持迭代访问（for 循环/next()）
适用场景	中小规模数据集，需随机访问、快速预处理	超大规模数据集，无法全量下载/存储

2. Dataset 核心操作

（1）索引访问

• 按行索引：返回单条数据字典

  dataset[0]  # 第一条数据
  dataset[-1] # 最后一条数据
  

• 按列索引：返回整列数据列表

  dataset["text"] # 所有文本数据
  

• 行列组合索引：返回指定位置的字段值

  dataset[0]["text"] # 第一条数据的文本
  dataset["text"][0] # 等价写法，顺序不影响结果
  

（2）切片访问

使用 : 符号获取数据集子集，返回字典（键为列名，值为对应切片列表）：

dataset[:3]   # 前3条数据
dataset[3:6]  # 第4-6条数据

3. IterableDataset 核心操作

（1）迭代访问

• 方式1：for 循环遍历

  iterable_dataset = load_dataset("ethz/food101", split="train", streaming=True)
  for example in iterable_dataset:
      print(example)
      break
  

• 方式2：next() + iter() 获取单条数据

  next(iter(iterable_dataset))
  

（2）列索引访问

支持按列名索引，返回列数据的迭代器：

next(iter(iterable_dataset["label"]))

（3）获取子集

使用 take(n) 方法获取前 n 条数据，返回新的 IterableDataset：

subset = iterable_dataset.take(3)
list(subset) # 转换为列表查看

（4）转换关系

Dataset 可通过 to_iterable_dataset() 转换为 IterableDataset（基于本地文件流式读取，速度更快）。

四、数据集预处理（核心流程）

预处理的核心目标是将原始数据转换为模型可接受的输入格式，不同模态数据预处理流程不同。

1. 文本数据集：分词（Tokenization）

（1）核心步骤

1. 加载数据集与预训练分词器

   from transformers import AutoTokenizer
   tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
   dataset = load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
   

2. 定义分词函数，使用 map() 批量处理

   def tokenization(example):
       return tokenizer(example["text"], truncation=True, padding=True)
   # batched=True 启用批量处理，提升效率
   dataset = dataset.map(tokenization, batched=True)
   

3. 设置数据集格式，适配机器学习框架
   • 适配 PyTorch：使用 set_format()

     dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
     

   • 适配 TensorFlow：使用 to_tf_dataset() + 数据收集器

     from transformers import DataCollatorWithPadding
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="tf")
     tf_dataset = dataset.to_tf_dataset(
         columns=["input_ids", "attention_mask"],
         label_cols=["label"],
         batch_size=2,
         collate_fn=data_collator,
         shuffle=True
     )
     

（2）关键说明

• 分词器需与预训练模型匹配，保证文本拆分规则一致。
• map() 是数据集预处理的核心方法，支持批量处理、多进程加速。

2. 音频数据集：重采样（Resampling）

（1）核心步骤

1. 加载数据集、特征提取器与 Audio 特征

   from transformers import AutoFeatureExtractor
   from datasets import load_dataset, Audio
   feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base-960h")
   dataset = load_dataset("PolyAI/minds14", "en-US", split="train")
   

2. 统一采样率：使用 cast_column() 转换音频列采样率

   # 将数据集采样率转换为模型所需的 16kHz
   dataset = dataset.cast_column("audio", Audio(sampling_rate=16_000))
   

3. 定义预处理函数，批量提取音频特征

   def preprocess_function(examples):
       audio_arrays = [x["array"] for x in examples["audio"]]
       inputs = feature_extractor(
           audio_arrays, sampling_rate=feature_extractor.sampling_rate, max_length=16000, truncation=True
       )
       return inputs
   dataset = dataset.map(preprocess_function, batched=True)
   

（2）关键说明

• 采样率需与预训练模型一致，否则会严重影响模型性能。
• Audio 特征会自动解码音频文件，cast_column() 可修改解码参数（如采样率）。

3. 图像数据集：数据增强（Data Augmentation）

（1）核心步骤

1. 加载数据集、特征提取器与 Image 特征

   from transformers import AutoFeatureExtractor
   from datasets import load_dataset, Image
   feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
   dataset = load_dataset("AI-Lab-Makerere/beans", split="train")
   

2. 统一图像格式（如 RGB）

   dataset = dataset.cast_column("image", Image(mode="RGB"))
   

3. 定义数据增强流水线（以 Albumentations 为例）

   import albumentations as A
   import numpy as np
   from PIL import Image
   
   transform = A.Compose([
       A.RandomCrop(height=256, width=256, pad_if_needed=True, p=1),
       A.HorizontalFlip(p=0.5),
       A.ColorJitter(p=0.5)
   ])
   
   def albumentations_transforms(examples):
       transformed_images = []
       for image in examples["image"]:
           # PIL → NumPy 数组
           image_np = np.array(image.convert("RGB"))
           # 应用增强
           transformed_image = transform(image=image_np)["image"]
           # NumPy → PIL
           pil_image = Image.fromarray(transformed_image)
           transformed_images.append(pil_image)
       examples["pixel_values"] = transformed_images
       return examples
   

4. 应用增强：使用 with_transform() 动态应用

   dataset = dataset.with_transform(albumentations_transforms)
   

（2）关键说明

• 数据增强库（如 Albumentations）通常要求输入为 NumPy 数组，需在 PIL 与 NumPy 之间转换。
• with_transform() 是即时转换，访问数据时才执行增强，节省内存。

五、创建自定义数据集

1. 文件格式构建器（支持常见格式）

直接通过 load_dataset() 加载本地文件，支持 csv/json/parquet/txt 等格式：

# 加载单个 CSV 文件
dataset = load_dataset("csv", data_files="my_file.csv")
# 加载多个文件（传入列表）
dataset = load_dataset("csv", data_files=["train.csv", "test.csv"])
# 按拆分加载（传入字典）
dataset = load_dataset("csv", data_files={"train": "train.csv", "test": "test.csv"})

2. 文件夹构建器（低代码快速创建图像/音频数据集）

适用于图像、音频数据集，自动根据文件夹结构生成拆分和标签。

（1）文件夹结构要求

dataset_root/
├── train/
│   ├── class_A/  # 标签 A 的样本
│   └── class_B/  # 标签 B 的样本
└── test/
    ├── class_A/
    └── class_B/

• 拆分由一级目录名（如 train/test）决定。
• 标签由二级目录名（如 class_A/class_B）决定。

（2）加载方法

• 图像数据集：使用 imagefolder

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

• 音频数据集：使用 audiofolder

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

（3）添加元数据

通过 metadata.csv 文件添加额外信息（如文本描述、转录文本），文件需包含 file_name 列关联样本文件：

file_name,text
sample1.png,a photo of class A
sample2.png,a photo of class B

3. 从 Python 数据结构创建

（1）从字典创建（from_dict()）

适用于小规模数据，直接将字典转换为 Dataset：

from datasets import Dataset
ds = Dataset.from_dict({
    "pokemon": ["bulbasaur", "squirtle"],
    "type": ["grass", "water"]
})

（2）从生成器创建（from_generator()）

内存高效，适用于大规模数据，逐行生成数据：

from datasets import Dataset
def gen():
    yield {"pokemon": "bulbasaur", "type": "grass"}
    yield {"pokemon": "squirtle", "type": "water"}
ds = Dataset.from_generator(gen)

（3）创建图像/音频数据集

通过 cast_column() 指定特征类型，自动解码文件：

from datasets import Dataset, Audio
audio_dataset = Dataset.from_dict({"audio": ["audio1.wav", "audio2.wav"]}).cast_column("audio", Audio())

六、核心总结

1. 加载流程：调研（load_dataset_builder）→ 加载（load_dataset，指定配置/拆分）→ 操作（索引/切片/迭代）→ 预处理（map/set_format）。
2. 预处理核心：文本分词、音频重采样、图像增强，最终转换为框架兼容格式。
3. 自定义数据集：小文件用格式构建器，图像/音频用文件夹构建器，Python 数据用 from_dict/from_generator。
4. 关键设计：内存映射、流式加载、批量处理，保证高效处理大规模数据集。

我可以帮你整理这份笔记的核心知识点速查表，方便你快速复习，需要吗？