本文旨在提供一种解决在使用 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 和其他参数，以达到最佳性能。
