使用 Transformers 解决 BERT 词嵌入中的内存溢出问题

本文旨在提供一种解决在使用 BERT 等 Transformers 模型进行词嵌入时遇到的内存溢出问题的有效方法。通过直接使用 tokenizer 处理文本输入，并适当调整 batch size，可以避免 `batch_encode_plus` 可能带来的内存压力，从而顺利生成词嵌入。

在使用 BERT 或其他 Transformers 模型生成文本数据集的词嵌入时，经常会遇到 OutOfMemoryError 错误，尤其是在处理长文本序列时。这主要是因为模型需要加载大量数据到 GPU 内存中进行计算。本文将介绍一种更高效的方法，通过优化文本处理流程和调整 batch size 来解决这个问题。

解决方案：优化文本处理和 Batch Size

传统的 batch_encode_plus 方法可能会导致内存占用过高。一种更有效的方法是直接使用 tokenizer 处理文本输入，并结合适当的 batch size。

1. 直接使用 Tokenizer 处理文本

不再使用 batch_encode_plus，而是直接使用 tokenizer 对象处理文本列表。这允许 tokenizer 内部更有效地管理内存。

import torch
from transformers import AutoModel, AutoTokenizer
# 输入文本列表 (可以是长句子)
texts = ['test1', 'test2']
# 加载预训练模型和 tokenizer
model_name = "indolem/indobert-base-uncased" # 这里替换为你想要使用的模型
model = AutoModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# 对文本进行分词、截断和填充
tokenized_texts = tokenizer(texts,
max_length=512,
truncation=True,
padding=True,
return_tensors='pt')

代码解释:

• AutoModel.from_pretrained(model_name): 加载指定名称的预训练模型。
• AutoTokenizer.from_pretrained(model_name): 加载与模型对应的 tokenizer。
• tokenizer(…): 使用 tokenizer 直接处理文本列表，设置最大长度、截断和填充策略，并返回 PyTorch 张量。

2. 前向传播

将 tokenizer 处理后的文本批次传递给模型进行前向传播。

# 前向传播
with torch.no_grad():
input_ids, attention_mask = tokenized_texts['input_ids'], tokenized_texts['attention_mask']
outputs = model(input_ids=input_ids,
attention_mask=attention_mask)
word_embeddings = outputs.last_hidden_state

代码解释:

• with torch.no_grad():: 禁用梯度计算，减少内存占用。
• outputs = model(…): 将输入 ID 和 attention mask 传递给模型进行前向传播。
• word_embeddings = outputs.last_hidden_state: 获取最后一层的隐藏状态，即词嵌入。

3. 检查输出形状

验证词嵌入的形状是否符合预期。

print(word_embeddings.shape)
# 输出: torch.Size([batch_size, num_seq_tokens, embed_size])
# 例如: torch.Size([2, 4, 768])

代码解释:

• word_embeddings.shape: 打印词嵌入的形状，通常为 [batch_size, num_seq_tokens, embed_size]，其中 batch_size 是批次大小，num_seq_tokens 是序列中的 token 数量，embed_size 是嵌入维度。

4. 调整 Batch Size (如果仍然出现 OOM)

如果即使使用上述方法仍然出现 OutOfMemoryError，则需要减小 batch size。可以循环处理数据，每次处理较小的批次。

import torch
from transformers import AutoModel, AutoTokenizer
# 输入文本列表
texts = ['test1', 'test2', 'test3', 'test4', 'test5']
# 加载预训练模型和 tokenizer
model_name = "indolem/indobert-base-uncased" # 这里替换为你想要使用的模型
model = AutoModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
batch_size = 2  # 设置较小的 batch size
for i in range(0, len(texts), batch_size):
batch_texts = texts[i:i + batch_size]
# 对文本进行分词、截断和填充
tokenized_texts = tokenizer(batch_texts,
max_length=512,
truncation=True,
padding=True,
return_tensors='pt')
# 前向传播
with torch.no_grad():
input_ids, attention_mask = tokenized_texts['input_ids'], tokenized_texts['attention_mask']
outputs = model(input_ids=input_ids,
attention_mask=attention_mask)
word_embeddings = outputs.last_hidden_state
print(f"Batch {i//batch_size + 1} embeddings shape: {word_embeddings.shape}")
# 在这里处理词嵌入，例如存储或进一步分析

代码解释:

• batch_size = 2: 设置较小的 batch size，例如 2。
• for i in range(0, len(texts), batch_size):: 循环处理数据，每次处理一个批次。
• batch_texts = texts[i:i + batch_size]: 提取当前批次的文本。
• 循环内部的代码与之前的示例相同，但现在每次处理的文本量较小，从而降低了内存占用。

注意事项

• 选择合适的 Batch Size: Batch size 的选择取决于 GPU 的内存大小和模型的复杂度。可以尝试不同的 batch size，找到一个既能充分利用 GPU 资源又能避免内存溢出的值。
• 使用 GPU: 确保代码在 GPU 上运行，这可以显著提高计算速度。
• 优化模型: 如果可能，可以尝试使用更小的模型或对模型进行量化，以减少内存占用。
• 梯度累积: 在某些情况下，可以使用梯度累积来模拟更大的 batch size，而无需增加内存占用。

总结

通过直接使用 tokenizer 处理文本输入并适当调整 batch size，可以有效地解决在使用 Transformers 模型进行词嵌入时遇到的内存溢出问题。这种方法简单易行，并且适用于各种 Transformers 模型。在实际应用中，可以根据具体情况调整 batch size 和其他参数，以达到最佳性能。