使用 Transformers 解决 BERT 词嵌入中的内存问题-Python教程-PHP中文网

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

霞舞

发布： 2025-10-19 12:55:11

原创

613人浏览过

使用 transformers 解决 bert 词嵌入中的内存问题

本文旨在解决在使用 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，即词嵌入。

AI建筑知识问答

用人工智能ChatGPT帮你解答所有建筑问题

查看详情

结果分析

获取的词嵌入的形状为 [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 的批次，并逐批处理。通过降低批次大小，可以显著减少内存占用。