本文旨在解决在使用 BERT 等 Transformer 模型进行词嵌入时遇到的内存不足问题。通过直接使用 tokenizer 处理文本输入，避免 `batch_encode_plus` 可能带来的问题。同时，提供了降低批次大小以进一步优化内存使用的建议，帮助用户高效地生成词嵌入。
在使用 BERT 或其他 Transformer 模型处理大量文本数据生成词嵌入时，OutOfMemoryError 是一个常见的问题。这通常是由于模型参数过多、输入序列过长或批次大小过大造成的。本文提供一种更高效的方法，通过优化文本预处理流程和调整批次大小来解决这个问题。
优化文本预处理
通常，我们会先使用 batch_encode_plus 对文本进行分词和编码，然后再将其输入到模型中。然而，对于长文本数据集，这种方法可能会导致内存占用过高。一种更优的方案是直接使用 tokenizer 处理文本输入，让 tokenizer 自身处理文本的截断、填充等操作。
以下是使用 AutoModel 和 AutoTokenizer 的示例代码：
import torch
from transformers import AutoModel, AutoTokenizer
# 输入文本列表 (可以是长句子)
texts = ['This is a test sentence.', 'Another test sentence.']
# 加载预训练模型和 tokenizer
model_name = "indolem/indobert-base-uncased" # 这里替换成你想要使用的模型
model = AutoModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# 使用 tokenizer 对文本进行分词、截断和填充
tokenized_texts = tokenizer(texts,
max_length=512, # 根据实际情况调整
truncation=True,
padding=True,
return_tensors='pt')这段代码首先加载了预训练的 BERT 模型和 tokenizer。然后，它使用 tokenizer 对文本进行分词、截断和填充，并将结果转换为 PyTorch 张量。通过这种方式，tokenizer 可以更好地管理内存，避免 batch_encode_plus 可能带来的问题。
模型前向传播
接下来，将编码后的文本输入到模型中进行前向传播，获取词嵌入：
# 前向传播
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这段代码使用 torch.no_grad() 上下文管理器禁用梯度计算，以减少内存占用。然后，它将 input_ids 和 attention_mask 输入到模型中，获取 last_hidden_state，即词嵌入。
结果分析
获取的词嵌入的形状为 [batch_size, num_seq_tokens, embed_size]，其中 batch_size 是批次大小，num_seq_tokens 是序列长度，embed_size 是嵌入维度。
print(word_embeddings.shape) # output: torch.Size([2, 4, 768])
进一步优化：降低批次大小
如果仍然遇到 OutOfMemoryError，可以尝试降低批次大小。这意味着将数据集分成更小的批次进行处理。
batch_size = 8 # 根据实际情况调整
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
# 对 word_embeddings 进行后续处理这段代码将数据集分成大小为 batch_size 的批次，并逐批处理。通过降低批次大小，可以显著减少内存占用。
总结与注意事项
通过直接使用 tokenizer 处理文本输入和降低批次大小，可以有效地解决在使用 BERT 等 Transformer 模型进行词嵌入时遇到的内存不足问题。
注意事项：
- max_length 参数需要根据数据集的实际情况进行调整。过小的 max_length 可能会导致信息丢失，过大的 max_length 会增加内存占用。
- 选择合适的预训练模型也很重要。较大的模型通常具有更好的性能，但也需要更多的内存。
- 如果仍然遇到内存问题，可以考虑使用更小的模型或增加 GPU 内存。
- 根据实际情况，可以尝试使用梯度累积等技术来进一步优化内存使用。
- 在Colab上使用GPU时，确保已经选择了GPU运行时环境。
- 可以尝试使用torch.cuda.empty_cache()释放不再使用的GPU内存。
通过以上方法，可以更有效地使用 Transformer 模型生成词嵌入，并避免 OutOfMemoryError。
