Hugging Face AutoTokenizer大型数据集内存优化指南

2025-02-28 13:05:57

Hugging Face AutoTokenizer 处理大型数据集的内存瓶颈问题

使用 Hugging Face 的 AutoTokenizer 对大型数据集进行 token 化处理时，经常遇到内存瓶颈，即便用了 batch_encode_plus 并进行截断，进程还是会时不时崩溃。这篇博客就是为了解决这个问题。

问题原因

AutoTokenizer 默认会将所有文本加载到内存中进行处理。当数据集过大时，这种方式很容易导致内存溢出。batch_encode_plus 允许分批处理,但如果批次大小设置不当，或者单个文本过长, 仍然可能耗尽内存。另外，某些 tokenizer (例如 SentencePiece) 内部实现可能导致在处理特定文本时产生较大的内存开销。

解决方案

针对这个问题，我们可以采用多种策略来降低内存占用。下面我将详细解释几种方案。

1. 减小批次大小 (Reduce Batch Size)

这是最直接的方法。减小 batch_encode_plus 方法中的 batch_size 参数可以降低每次处理的文本数量，从而降低内存峰值。

原理：

较小的批次意味着每次加载到内存中的数据更少。

代码示例：

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
texts = ["Example text 1", "Example text 2", ..., "Example text N"] # 你的大型文本列表
encoded_inputs = tokenizer.batch_encode_plus(
    texts,
    padding=True,
    truncation=True,
    max_length=512,
    return_tensors="pt",
    batch_size=16  # 原来可能是 32 或 64，现在减小到 16 或更小
)

注意事项：

过小的批次大小可能会降低 tokenization 速度，需要在速度和内存占用之间找到平衡。可以尝试不同的批次大小，例如 8, 4, 甚至 1，观察内存使用情况和处理速度。

2. 使用 `Dataset.map` 的 `batched=True`

Hugging Face datasets 库的 Dataset.map 方法提供了一个 batched 参数，可以进行批处理,并且更高效。

原理：

Dataset.map 会按批次读取数据, 对批次进行 tokenization, 然后将结果写回, 避免将整个数据集加载到内存中. batched=True 会把数据分成小批处理，并且每次处理一个小批。

代码示例：

from datasets import load_dataset
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# 假设你有一个大型数据集，例如 "imdb"
dataset = load_dataset("imdb", split="train")

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

# 使用 batched=True 进行批处理 tokenization
tokenized_datasets = dataset.map(tokenize_function, batched=True, batch_size=1000) #调整这个batch_size

#如果担心cache文件导致OOM，增加remove_columns 移除原始文本
#tokenized_datasets = dataset.map(tokenize_function, batched=True, batch_size=1000,remove_columns=["text"])

进阶技巧

调整 batch_size 参数：根据你的机器内存和数据集特性调整, 从小到大尝试（例如 100，500，1000,2000）。
使用 num_proc 参数： map 方法还支持多进程处理 (num_proc)，可以进一步加快处理速度，前提是你机器有多余CPU核心。

3. 流式处理 (Streaming)

如果数据集实在太大, 无法全部加载, 可以使用 datasets 库的流式处理功能.

原理：

流式处理每次只从磁盘或网络读取一小部分数据，进行 tokenization, 处理完后再读取下一部分。

代码示例：

from datasets import load_dataset
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# 使用 streaming=True 加载数据集
dataset = load_dataset("imdb", split="train", streaming=True)

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

tokenized_datasets = dataset.map(tokenize_function, batched=True)
#处理流数据集的批处理
#请注意，在使用流数据集时，批大小处理可能会有所不同。
# 流数据集通常需要更小的批大小，或使用不同的迭代策略。

#迭代流式处理
for example in tokenized_datasets:
     # 对单个 example 进行处理
    pass

注意： 流式处理模式下，有些操作(例如 shuffle) 不支持或需要特别处理. 并且, 对 map 之后的 dataset 进行操作时要特别注意，不要触发将整个数据集加载到内存的操作.

4. 使用磁盘映射 (Disk Mapping/Memory Mapping)

某些情况下, 可以使用 datasets 的磁盘映射功能 (如果数据集格式支持) 将数据集映射到磁盘上的文件, 而不是全部加载到内存.

原理:

操作系统级别的内存映射技术, 允许你像访问内存一样访问文件, 但实际数据按需从磁盘加载。

代码示例：

from datasets import load_dataset
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# 加载数据集，自动使用 memory mapping (如果数据集格式支持，例如 arrow 格式)
dataset = load_dataset("your_dataset_name", split="train") #需要先处理成datasets需要的格式

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

注意 : 这个特性依赖于数据集格式和底层实现，不一定所有数据集都支持。.arrow文件是推荐支持的格式。

5. 手动分块处理 (Manual Chunking)

如果以上方法都不奏效, 你可以手动将数据集分割成多个小文件, 逐个处理。

原理：

将一个大文件分解成多个小文件，每次只处理一个小文件，避免一次性加载整个数据集。

步骤：

将原始数据集分割成多个小文件 (例如，每个文件包含 10000 条数据)。
编写一个循环，逐个读取这些小文件。
对每个小文件的数据进行 tokenization，并将结果保存。
（可选）最后将所有 tokenization 后的结果合并。

代码示例：

# 假设你已经将数据集分割成多个文件：data_part_1.txt, data_part_2.txt, ...

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def tokenize_and_save(input_file, output_file):
    with open(input_file, "r", encoding="utf-8") as f:
        texts = f.readlines()

    encoded_inputs = tokenizer.batch_encode_plus(
        texts,
        padding=True,
        truncation=True,
        max_length=512,
        return_tensors="pt",
        batch_size=32  # 根据实际情况调整
    )

    # 将 encoded_inputs 保存到 output_file (例如，使用 torch.save)
    # ...

file_list = ["data_part_1.txt", "data_part_2.txt", ...] #所有小文件名字
for i, file in enumerate(file_list):
    tokenize_and_save(file, f"output_part_{i}.pt")

6. 更换 Tokenizer (更换为更节省内存的)

如果特定 tokenizer 内存占用过高，可以考虑更换为更节省内存的 tokenizer。

原理:

不同的 tokenizer 内部实现不同，内存占用可能有差异。一些 tokenizer (例如，基于 SentencePiece 的) 可能比 WordPiece tokenizer 更节省内存。

举例：

如果你原来使用 bert-base-uncased, 可以尝试更换为 distilbert-base-uncased 或 albert-base-v2 的 tokenizer. 它们通常在保持性能的同时, 有更小的模型和更低的内存占用.

代码示例：

from transformers import AutoTokenizer

# 原来：
# tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# 尝试更换：
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
# 或者
tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")

7.数据过滤与采样

预处理中, 先过滤掉过长的文本, 或进行采样.

原理:
减少长文本的数量可以直接降低内存需求, 因为长文本通常是造成内存问题的元凶。
数据采样可以在保持数据集分布的同时减少总体数据量。

代码示例（数据过滤）：

from datasets import load_dataset

dataset = load_dataset("imdb", split="train")

# 过滤掉长度超过 1024 的文本
filtered_dataset = dataset.filter(lambda example: len(example["text"].split()) <= 1024)

代码示例（数据采样）:

from datasets import load_dataset

dataset = load_dataset("imdb", split="train")

# 从数据集中随机采样 10% 的数据
sampled_dataset = dataset.shuffle(seed=42).select(range(int(len(dataset)*0.1)))

综合应用:
处理大型数据集时，最佳方案往往是多种方法的组合. 比如，先使用 Dataset.map 的 batched=True 进行初步处理, 如果仍然有问题，再结合减小批次大小, 数据过滤等方法. 根据数据集特性和硬件资源,灵活调整.

希望上面方法可以帮你搞定大文本的tokenizer内存溢出问题。