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

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

心靈之曲

发布： 2025-10-18 15:35:18

原创

179人浏览过

使用 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')

登录后复制

代码解释:

AI建筑知识问答

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

查看详情

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 和其他参数，以达到最佳性能。

以上就是使用 Transformers 解决 BERT 词嵌入中的内存溢出问题的详细内容，更多请关注php中文网其它相关文章！