当前位置：首页 > news >正文

Transformers Dataset 完全指南：从基础到高级应用

news 2025/9/27 12:34:51

Transformers Dataset 完全指南：从基础到高级应用

本教程将全面介绍 Hugging Face Transformers 库中的数据集处理，涵盖从基础概念到高级应用的各个方面。

1. 数据集基础概念

1.1 什么是 Dataset？

在 Transformers 中，Dataset 是一个高效的数据容器，专门为机器学习任务设计，具有以下特点：

内存映射：支持大型数据集的高效加载
流式处理：可以处理超过内存大小的数据集
格式统一：提供标准化的接口
预处理管道：内置数据转换功能

1.2 安装和环境设置

# 安装必要的库
uv init
uv add 'transformers[torch]'
uv add 'datasets' 
uv add 'jupyter'# 导入常用模块
from datasets import Dataset, DatasetDict, load_dataset, load_from_disk
import pandas as pd
import numpy as np
from transformers import AutoTokenizer, AutoModel

2. 内置数据集的使用

2.1 加载 Hugging Face 数据集

# 加载公开数据集
dataset = load_dataset("imdb")# 查看数据集结构
print(dataset)
print(dataset["train"][0])  # 查看第一条训练数据
print(dataset["train"].features)  # 查看第一条训练数据
print(dataset["train"].num_rows)  # 查看第一条训练数据
# 加载特定子集
dataset = load_dataset("glue", "mrpc")  # MRPC 任务
dataset = load_dataset("squad", split="train")  # 只加载训练集

输出结果：

DatasetDict({train: Dataset({features: ['text', 'label'],num_rows: 25000})test: Dataset({features: ['text', 'label'],num_rows: 25000})unsupervised: Dataset({features: ['text', 'label'],num_rows: 50000})
})
{'text': 'I rented I AM CURIOUS-YELLOW from my video store because of all the controversy that surrounded it when it was first released in 1967. I also heard that at first it was seized by U.S. customs if it ever tried to enter this country, therefore being a fan of films considered "controversial" I really had to see this for myself.<br /><br />The plot is centered around a young Swedish drama student named Lena who wants to learn everything she can about life. In particular she wants to focus her attentions to making some sort of documentary on what the average Swede thought about certain political issues such as the Vietnam War and race issues in the United States. In between asking politicians and ordinary denizens of Stockholm about their opinions on politics, she has sex with her drama teacher, classmates, and married men.<br /><br />What kills me about I AM CURIOUS-YELLOW is that 40 years ago, this was considered pornographic. Really, the sex and nudity scenes are few and far between, even then it\'s not shot like some cheaply made porno. While my countrymen mind find it shocking, in reality sex and nudity are a major staple in Swedish cinema. Even Ingmar Bergman, arguably their answer to good old boy John Ford, had sex scenes in his films.<br /><br />I do commend the filmmakers for the fact that any sex shown in the film is shown for artistic purposes rather than just to shock people and make money to be shown in pornographic theaters in America. I AM CURIOUS-YELLOW is a good film for anyone wanting to study the meat and potatoes (no pun intended) of Swedish cinema. But really, this film doesn\'t have much of a plot.', 'label': 0}
{'text': Value('string'), 'label': ClassLabel(names=['neg', 'pos'])}
25000

Dataset****（数据集实例）：
是单个数据集的具体载体，存储单一子集的原始数据（如训练集、测试集等）。
从你的示例来看，train、test、unsupervised都是Dataset对象，每个都包含：
- features：定义数据的字段结构（如['text', 'label']，即每条数据有 “文本” 和 “标签” 两个字段）；
- num_rows：该子集的样本数量（如 train 有 25000 条，unsupervised 有 50000 条）。

简单说，Dataset就是 “具体的数据容器”，直接对应某一个拆分的数据集（如训练集单独的所有样本）。

DatasetDict****（数据集字典）：
是**多个****Dataset**的集合容器，本质上是一个 “字典 - like” 的结构，用于管理同一数据集的不同拆分（split）。
你的示例中，DatasetDict包含了 3 个键值对：
- 键（key）：拆分名称（如'train'、'test'、'unsupervised'）；
- 值（value）：对应的Dataset对象（即每个拆分的具体数据）。

简单说，DatasetDict是 “拆分的管理者”，将同一来源的不同子集（训练 / 测试 / 无监督等）整合在一起，方便统一操作。

2.2 数据集分割与查看

# 数据集通常包含多个分割
dataset = load_dataset("imdb")
print(f"分割: {list(dataset.keys())}")
print(f"训练集大小: {len(dataset['train'])}")
print(f"测试集大小: {len(dataset['test'])}")# 查看数据格式
sample = dataset["train"][0]
print(f"样本键: {list(sample.keys())}")
print(f"文本: {sample['text']}")
print(f"标签: {sample['label']}")

2.3 数据集切片和采样

# 选择前1000个样本
small_dataset = dataset["train"].select(range(1000))# 随机采样
sampled_dataset = dataset["train"].shuffle(seed=42).select(range(500))# 过滤数据
def filter_long_reviews(example):return len(example["text"].split()) < 200  # 只保留短评论文本filtered_dataset = dataset["train"].filter(filter_long_reviews)
print(f"过滤后大小: {len(filtered_dataset)}")

3. 自定义数据集

3.1 从 Pandas DataFrame 创建

# 创建示例数据
data = {"text": ["I love this movie, it's amazing!","This is the worst film I've ever seen.","The acting was good but the plot was weak.","Absolutely fantastic cinematography!"],"label": [1, 0, 0, 1],"source": ["user1", "user2", "user3", "user4"]
}df = pd.DataFrame(data)
custom_dataset = Dataset.from_pandas(df)
print(custom_dataset)

输出：

Dataset({features: ['text', 'label', 'source'],num_rows: 4
})

3.2 从字典创建

# 直接从字典创建
data_dict = {"text": ["Text 1", "Text 2", "Text 3"],"label": [0, 1, 0],"id": [1, 2, 3]
}dataset_from_dict = Dataset.from_dict(data_dict)
print(dataset_from_dict)

输出：

Dataset({features: ['text', 'label', 'id'],num_rows: 3
})

3.3 从文件创建

# 从 CSV 文件
dataset_csv = Dataset.from_csv("path/to/your/data.csv")# 从 JSON 文件
dataset_json = Dataset.from_json("path/to/your/data.json")# 从文本文件（每行一个样本）
def create_text_dataset(file_path):with open(file_path, 'r', encoding='utf-8') as f:lines = f.readlines()data = {"text": [line.strip() for line in lines]}return Dataset.from_dict(data)text_dataset = create_text_dataset("sample_texts.txt")

CSV 文件格式（对应 Dataset.from_csv）

CSV（逗号分隔值）文件需包含表头（字段名），每行代表一条样本，字段间用逗号分隔。
假设你的数据包含 text（文本内容）和 label（标签）两个字段，示例 data.csv 如下：

text,label
"这是第一条文本，用于测试",0
"第二条样本内容，长度可以不同",1
"CSV文件用逗号分隔字段，若文本含逗号需用引号包裹",0
"最后一条数据",1

注意：如果文本中包含逗号、换行符等特殊字符，需用双引号 "" 将文本字段包裹（如第三条数据）。
表头必须与代码中后续处理的字段名一致（如 example["text"] 对应表头 text）。

JSON 文件格式（对应 Dataset.from_json）

JSON 文件需为列表形式，列表中的每个元素是一个字典，代表一条样本，字典的键为字段名。
示例 data.json 如下（与 CSV 内容对应）：

[{"text": "这是第一条文本，用于测试", "label": 0},{"text": "第二条样本内容，长度可以不同", "label": 1},{"text": "JSON文件用列表包裹，每条数据是一个字典", "label": 0},{"text": "最后一条数据", "label": 1}
]

必须是JSON 数组（外层用 [] 包裹），否则 from_json 会报错。
字段名（如 text、label）需保持一致，且避免语法错误（如逗号遗漏、引号不匹配）。

TXT 文件格式（对应自定义的 create_text_dataset）

TXT 文件采用每行一个样本的格式，每行文本会被作为 text 字段的值，无需表头或额外标记。
示例 sample_texts.txt 如下：

这是第一条文本，直接占一行
第二条样本，同样单独一行
TXT文件不需要表头，每行就是一个text字段的内容
空行会被处理为空白文本（如果需要过滤可在代码中添加strip()后判断）
最后一行文本

代码中 line.strip() 会去除每行首尾的空白字符（空格、换行符等），但保留行内的空格。
若存在空行，会被转换为 ""（空字符串），可在生成数据集后用 filter 过滤（如 filter(lambda x: x["text"] != "")）。

3.4 创建 DatasetDict

# 创建完整的数据集字典（包含训练集、验证集、测试集）
train_data = {"text": ["train1", "train2"], "label": [0, 1]}
val_data = {"text": ["val1", "val2"], "label": [1, 0]}
test_data = {"text": ["test1", "test2"], "label": [0, 1]}dataset_dict = DatasetDict({"train": Dataset.from_dict(train_data),"validation": Dataset.from_dict(val_data),"test": Dataset.from_dict(test_data)
})print(dataset_dict)

4. 数据预处理与特征工程

4.1 使用 Tokenizer

from datasets import Dataset, DatasetDict, load_dataset, load_from_diskimport pandas as pd
import numpy as np
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")# 基本 tokenization 函数
def tokenize_function(examples):return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=10)
dataset=Dataset.from_dict({"text":["hello world","this is a big eggs"],"label":[0,1]
})
# 应用 tokenization
tokenized_dataset = dataset.map(tokenize_function, batched=True,remove_columns=["text","label"], batch_size=500)
print(tokenized_dataset[0])
print(tokenized_dataset[1])
print(tokenizer.decode(tokenized_dataset[0]['input_ids']))

关于分词器相关概念参考:https://blog.csdn.net/liaomin416100569/article/details/139627429?spm=1001.2014.3001.5501 回顾预处理章节
map函数的参数，参考map章节。

输出:

{'input_ids': [101, 7592, 2088, 102, 0, 0, 0, 0, 0, 0], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 0, 0, 0, 0, 0, 0]}
{'input_ids': [101, 2023, 2003, 1037, 2502, 6763, 102, 0, 0, 0], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 0, 0, 0]}
[CLS] hello world [SEP] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD]

4.2 高级预处理管道

from datasets import Dataset, DatasetDict, load_dataset, load_from_diskimport pandas as pd
import numpy as np
from transformers import AutoTokenizer, AutoModeltokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")def preprocess_function(examples):# 文本清理texts = [text.lower().strip() for text in examples["text"]]print(texts)# Tokenizationtokenized = tokenizer(texts,padding="max_length",truncation=True,max_length=256,return_tensors="np"  # 返回 numpy 数组)print(examples["label"])# 添加标签tokenized["labels"] = examples["label"]return tokenizeddataset = DatasetDict({"train": Dataset.from_dict({"text": ["hello world", "this is a big eggs"],"label": [0, 1]}),"test": Dataset.from_dict({"text": ["test world", "this is a test world"],"label": [0, 1]})
})
# 应用预处理
processed_dataset = dataset.map(preprocess_function,batched=True,remove_columns=dataset["train"].column_names  # 移除原始列
)
print(processed_dataset)
print(processed_dataset["train"][0])

如果数据集时DatasetDict会将train作为一批传入processed_dataset一次，test作为一批传入processed_dataset一次

4.3 动态填充和批处理

from datasets import Dataset
from torch.utils.data import DataLoader
from transformers import AutoTokenizertokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")def tokenize_without_padding(examples):"""truncation=True：截断过长文本padding=True：填充过短文本例如 tokenizer(..., max_length=128, truncation=True, padding="max_length")，此时所有文本会被统一处理为 128 长度（过长截断，过短填充至 128）"""return tokenizer(examples["text"], truncation=True)dataset = Dataset.from_dict({"text": ["hello world", "this is a big eggs", "test world", "this is a test world"],"label": [0, 1, 0, 1]
})
# 先进行基础 tokenization
tokenized_dataset = dataset.map(tokenize_without_padding, batched=True)
# 使用 DataCollator 进行动态填充
from transformers import DataCollatorWithPadding#  准备分词后的数据集（转换为 PyTorch 张量格式）
train_dataset = tokenized_dataset.with_format("torch", columns=["input_ids", "attention_mask", "label"])data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
# 创建 DataLoader，用 data_collator 作为批处理函数
train_dataloader = DataLoader(train_dataset,batch_size=2,  # 批次大小为 2collate_fn=data_collator  # 用动态填充函数处理批次
)
# 查看处理后的批次数据,两个一循环
for batch in train_dataloader:print("批次 input_ids 形状：", batch["input_ids"].shape, "input_ids：", batch["input_ids"], "attention_mask：",batch["attention_mask"])  # 输出 (2, 6)，因为该批次最长为 6

输出:

批次 input_ids 形状： torch.Size([2, 7]) input_ids： tensor([[ 101, 7592, 2088,  102,    0,    0,    0],[ 101, 2023, 2003, 1037, 2502, 6763,  102]]) attention_mask： tensor([[1, 1, 1, 1, 0, 0, 0],[1, 1, 1, 1, 1, 1, 1]])
批次 input_ids 形状： torch.Size([2, 7]) input_ids： tensor([[ 101, 3231, 2088,  102,    0,    0,    0],[ 101, 2023, 2003, 1037, 3231, 2088,  102]]) attention_mask： tensor([[1, 1, 1, 1, 0, 0, 0],[1, 1, 1, 1, 1, 1, 1]])

当你用 map() 对文本进行分词时（如代码中的 tokenize_without_padding），如果没有设置 padding="max_length"，那么每条文本会保留原始分词后的长度（不同样本长度可能不同）。而模型输入要求同一批次的样本长度必须一致，此时就需要通过 DataCollatorWithPadding 动态调整：

动态确定批次最大长度：对每个批次的样本，自动找到该批次中最长的文本长度，作为当前批次的统一长度。
按需填充：只对该批次中短于最大长度的样本进行填充（用分词器的 pad_token，如 BERT 的 [PAD]），避免全局固定长度（如 max_length=128）导致的冗余填充。
自动处理注意力掩码（attention_mask）：填充部分对应的 attention_mask 会设为 0，模型会忽略这些位置，不影响计算。

对比两种填充方式的优缺点：

方式	做法	缺点	`DataCollatorWithPadding` 优势
静态填充（`map` 时）	`padding="max_length"` 固定长度	若全局设长了，会导致大量冗余填充，浪费计算资源；设短了可能截断过多文本。	只按批次最大长度填充，减少冗余，节省显存和计算时间。

5. 数据加载与批处理

5.1 创建 DataLoader

from torch.utils.data import DataLoader# 基本 DataLoader
dataloader = DataLoader(tokenized_dataset["train"],batch_size=32,shuffle=True,collate_fn=data_collator
)# 测试 DataLoader
for batch in dataloader:print(f"Batch keys: {batch.keys()}")print(f"Input shape: {batch['input_ids'].shape}")break  # 只看第一个批次

关键参数说明

参数	作用	建议值
`dataset`	输入的数据集（如 `tokenized_dataset["train"]`）	必须指定
`batch_size`	每批次样本数	常用 16/32/64（根据显存调整）
`shuffle`	是否打乱数据顺序	训练集 `True`，验证集 `False`
`collate_fn`	批次数据整理函数（如 `DataCollatorWithPadding`）	文本数据必用，处理填充等
`num_workers`	并行加载的进程数	通常设为 CPU 核心数的 1-2 倍

5.2 流式处理大型数据集

# 流式加载大型数据集
streaming_dataset = load_dataset("wikitext", "wikitext-103-v1", streaming=True)# 只加载部分数据
for i, example in enumerate(streaming_dataset["train"]):if i >= 5:  # 只查看前5个样本breakprint(example["text"])# 流式处理函数
def process_streaming_data(dataset, tokenizer, batch_size=1000):def process_batch(batch):return tokenizer(batch["text"], truncation=True, padding=True)return dataset.map(process_batch,batched=True,batch_size=batch_size)

5.3 分布式训练数据加载

from transformers import Trainer, TrainingArgumentstraining_args = TrainingArguments(output_dir="./results",per_device_train_batch_size=16,per_device_eval_batch_size=16,dataloader_num_workers=4,  # 多进程加载dataloader_pin_memory=True,  # 固定内存，加速GPU传输
)trainer = Trainer(model=model,args=training_args,train_dataset=tokenized_dataset["train"],eval_dataset=tokenized_dataset["test"],data_collator=data_collator,
)

6. 高级技巧与最佳实践

6.1 数据增强

# 文本数据增强
import nlpaug.augmenter.word as naw# 创建增强器（需要安装 nlpaug: pip install nlpaug）
aug = naw.ContextualWordEmbsAug(model_path='bert-base-uncased', action="insert")def augment_text(example):augmented_text = aug.augment(example["text"])return {"text": augmented_text, "label": example["label"]}# 应用数据增强（谨慎使用，可能很慢）
augmented_dataset = dataset.map(augment_text)# 更简单的方法：同义词替换
def synonym_replacement(example):words = example["text"].split()if len(words) > 1:# 简单的同义词替换逻辑import randomidx = random.randint(0, len(words)-1)words[idx] = f"_{words[idx]}_"  # 模拟替换return {"text": " ".join(words), "label": example["label"]}simple_augmented = dataset.map(synonym_replacement)

6.2 缓存机制

# 使用缓存加速预处理
cached_dataset = dataset.map(tokenize_function,batched=True,load_from_cache_file=True,  # 从缓存加载cache_file_name="./cache/tokenized_cache"  # 指定缓存文件
)# 检查缓存状态
print(f"缓存文件: {cached_dataset.cache_files}")

6.3 格式转换

# 转换为 PyTorch 张量格式
pt_dataset = tokenized_dataset.with_format("torch")
print(f"格式: {pt_dataset.format}")# 转换为 TensorFlow 张量格式
tf_dataset = tokenized_dataset.with_format("tensorflow")# 转换为 numpy 格式
np_dataset = tokenized_dataset.with_format("numpy")

6.4 数据集保存与加载

# 保存处理后的数据集
tokenized_dataset.save_to_disk("./processed_imdb")# 从磁盘加载
loaded_dataset = load_from_disk("./processed_imdb")# 导出为其他格式
tokenized_dataset.to_csv("./dataset.csv")
tokenized_dataset.to_json("./dataset.json")
tokenized_dataset.to_pandas().to_csv("./pandas_dataset.csv")

7. 实战案例

7.1 文本分类完整流程

from transformers import AutoTokenizer, AutoModelForSequenceClassification
from transformers import Trainer, TrainingArguments
import evaluate# 加载数据和模型
dataset = load_dataset("imdb")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)# 预处理函数
def preprocess_function(examples):return tokenizer(examples["text"], truncation=True, padding=True, max_length=256)# 应用预处理
tokenized_dataset = dataset.map(preprocess_function, batched=True)# 数据整理器
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)# 评估指标
metric = evaluate.load("accuracy")def compute_metrics(eval_pred):logits, labels = eval_predpredictions = np.argmax(logits, axis=-1)return metric.compute(predictions=predictions, references=labels)# 训练参数
training_args = TrainingArguments(output_dir="./results",learning_rate=2e-5,per_device_train_batch_size=16,per_device_eval_batch_size=16,num_train_epochs=3,weight_decay=0.01,evaluation_strategy="epoch",save_strategy="epoch",load_best_model_at_end=True,
)# 创建训练器
trainer = Trainer(model=model,args=training_args,train_dataset=tokenized_dataset["train"],eval_dataset=tokenized_dataset["test"],tokenizer=tokenizer,data_collator=data_collator,compute_metrics=compute_metrics,
)# 开始训练
trainer.train()

7.2 问答系统数据处理

# 加载 SQuAD 数据集
squad_dataset = load_dataset("squad")print("SQuAD 样本示例:")
print(squad_dataset["train"][0])# 预处理问答数据
def preprocess_qa_function(examples):questions = [q.strip() for q in examples["question"]]# 处理上下文和答案inputs = tokenizer(questions,examples["context"],max_length=384,truncation="only_second",stride=128,return_overflowing_tokens=True,return_offsets_mapping=True,padding="max_length",)# 处理答案起始位置offset_mapping = inputs.pop("offset_mapping")sample_map = inputs.pop("overflow_to_sample_mapping")answers = examples["answers"]start_positions = []end_positions = []for i, offset in enumerate(offset_mapping):sample_idx = sample_map[i]answer = answers[sample_idx]start_char = answer["answer_start"][0]end_char = start_char + len(answer["text"][0])# 查找 token 位置sequence_ids = inputs.sequence_ids(i)idx = 0while idx < len(sequence_ids) and sequence_ids[idx] != 1:idx += 1context_start = idxwhile idx < len(sequence_ids) and sequence_ids[idx] == 1:idx += 1context_end = idx - 1# 如果答案不在截取的上下文中，将位置设为 0if offset[context_start][0] > end_char or offset[context_end][1] < start_char:start_positions.append(0)end_positions.append(0)else:# 否则找到答案的 token 位置idx = context_startwhile idx <= context_end and offset[idx][0] <= start_char:idx += 1start_positions.append(idx - 1)idx = context_endwhile idx >= context_start and offset[idx][1] >= end_char:idx -= 1end_positions.append(idx + 1)inputs["start_positions"] = start_positionsinputs["end_positions"] = end_positionsreturn inputs# 应用预处理
tokenized_squad = squad_dataset.map(preprocess_qa_function, batched=True, remove_columns=squad_dataset["train"].column_names)

7.3 多模态数据处理

# 图像-文本多模态示例
from transformers import AutoImageProcessordef process_multimodal_data(examples):# 处理图像image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")images = [Image.open(path) for path in examples["image_path"]]image_inputs = image_processor(images, return_tensors="pt")# 处理文本text_inputs = tokenizer(examples["caption"],padding="max_length",truncation=True,max_length=128)# 合并输入return {**image_inputs, **text_inputs, "labels": examples["label"]}# 假设我们有包含图像路径和标题的数据集
multimodal_dataset = Dataset.from_dict({"image_path": ["img1.jpg", "img2.jpg"],"caption": ["A beautiful sunset", "A cute dog"],"label": [1, 0]
})processed_multimodal = multimodal_dataset.map(process_multimodal_data, batched=True)

7.4 自定义迭代器

class CustomDataIterator:def __init__(self, dataset, batch_size=32, shuffle=True):self.dataset = datasetself.batch_size = batch_sizeself.shuffle = shuffleself.indices = list(range(len(dataset)))if shuffle:np.random.shuffle(self.indices)self.current_idx = 0def __iter__(self):return selfdef __next__(self):if self.current_idx >= len(self.dataset):if self.shuffle:np.random.shuffle(self.indices)self.current_idx = 0raise StopIterationbatch_indices = self.indices[self.current_idx:self.current_idx + self.batch_size]batch = self.dataset[batch_indices]self.current_idx += self.batch_sizereturn batch# 使用自定义迭代器
custom_iterator = CustomDataIterator(tokenized_dataset["train"], batch_size=16)
for batch in custom_iterator:print(f"处理批次大小: {len(batch['input_ids'])}")