BERT推理显存爆满？7个实用技巧教你快速优化！显存优化的最佳实践指南

大多_C

已于 2024-09-20 11:01:38 修改

阅读量1.1k

点赞数 21

文章标签： bert 人工智能深度学习

于 2024-09-20 10:27:43 首次发布

本文链接：https://blog.csdn.net/weixin_46933702/article/details/142377737

版权

如果只使用BERT进行推理得到词向量，但显存仍然爆满，以下几个建议可以帮助缓解显存问题：

分批处理（Batching）：
即使是在推理阶段，也可以将输入数据分成较小的批次（batch），这样每次推理只处理较少的数据，从而减少显存占用。具体来说可以减小 batch size，逐步降低直到显存能够容纳。
减少输入序列长度：
BERT 的显存占用与输入序列的长度直接相关。可以尝试裁剪输入文本，使得序列长度减少。特别是对很长的文本，可以设置合理的最大长度（如128或256），从而减少计算需求。
使用梯度检查点（Gradient Checkpointing）：
虽然主要用于训练，但在推理时可以通过减少内存开销来节省显存。可以通过调整模型以支持这种技术，减少中间计算结果的缓存。
半精度推理（FP16）：
如果你的硬件支持，可以使用半精度浮点数（FP16）进行推理。许多深度学习框架（如 PyTorch）支持将模型加载为半精度，从而减少显存需求。这在大多数推理场景下对精度影响不大。
冻结部分层：
如果只需要较浅层的词向量，可以考虑冻结模型的某些层，只对前几层进行推理。例如，如果只需要基础的词向量特征，可以考虑只使用 BERT 的前几层，而不是整个模型。
使用更小的模型：
如果推理性能和速度比模型大小更重要，使用更小版本的 BERT，例如 DistilBERT 或 ALBERT，可以大幅减少显存占用，同时还能保持不错的性能。
迁移到更大显存的设备：
如果硬件允许，尝试迁移到具有更多显存的GPU。对于BERT这样的模型，如果输入数据量大且显存瓶颈严重，可能需要考虑硬件升级。

如果这些方法还无法解决问题，可以分享你具体的代码片段，我可以帮你进一步优化和调整。

在运行代码时，显存逐步增加的情况，可能是因为在 for 循环中处理句子时，每次的计算结果并未及时释放，导致显存占用累积。这里有几个常见的问题及优化建议，可以帮助你减少显存占用，并防止显存泄露。

1. 确保不保留计算图：

在推理过程中，如果没有进行反向传播，则不需要保留计算图。因此，你可以在获取模型输出后使用 torch.no_grad()，以确保不会存储反向传播所需的中间计算结果，这样可以大幅减少显存占用。

优化代码如下：

def get_entity_encodings(sentence):
    model.eval()  # 设置模型为评估模式
    with torch.no_grad():  # 不需要计算梯度
        ...     
        return e1_encoding, e2_encoding, tokens

这里通过 torch.no_grad() 关闭了梯度计算，不会为推理计算生成反向传播的计算图。

2. 释放显存（CUDA 清理）：

在每次处理句子后，可以调用 torch.cuda.empty_cache() 来手动清理 CUDA 中未使用的显存，防止显存逐步累积。它不会影响当前运行的代码，但可以帮助释放不再需要的显存。

在循环中可以定期清理显存：

for idx, (sentence, label) in enumerate(tqdm(zip(train_sentence, train_sentence_label), total=len(train_sentence), desc="Processing sentences")):
    try:
        ....
        # 定期释放 CUDA 显存
        if idx % 100 == 0:
            torch.cuda.empty_cache()
    except ValueError as e:
        print(f"跳过句子 {idx}: {e}")

3. Batching（批处理）：

虽然你现在是在逐句处理，但可以考虑使用批处理来一次处理多个句子，减少模型前向传播的次数，这样能更有效地使用显存并加快处理速度。批处理的实现方式如下：

使用 DataLoader 进行批次处理，调整 batch_size 以找到适合显存的大小。

from torch.utils.data import DataLoader

dataset = TensorDataset(inputs, sentence_labels, sentence_indices)
dataloader = DataLoader(dataset, batch_size=16, shuffle=True)  # 调整 batch_size 以适应显存

for batch in tqdm(dataloader, desc="Processing batches"):
    batch_inputs, batch_labels, batch_indices = batch
    # 批处理后的推理操作

4. 减少中间结果的存储：

目前的代码中，你存储了每一对实体编码。如果数据量较大，可以考虑在需要时才使用中间结果，而不是一直保存所有数据。例如，将这些中间结果保存到磁盘，或仅保存你确实需要的部分。

5. 使用更小的模型：

如果显存依然不够，可以尝试使用更小的模型，如 distilBERT 或 TinyBERT，它们的显存占用比 BERT-base 要小很多，且推理速度更快。

通过这些优化策略，你可以大幅减少推理过程中的显存占用，并防止显存泄漏导致的崩溃。

从图表中可以看出，尽管你使用了 torch.no_grad() 和 torch.cuda.empty_cache()，显存的占用仍然在逐步增加。这表明某些部分可能存在显存的积累，没有被正确释放。

可能的原因与解决方法：

BERT 模型输出的累积：
即使使用了 torch.no_grad()，model(**inputs) 仍然会返回 last_hidden_state，这些可能会在每次循环时被保留在 GPU 中。因为你将 e1_encodings 和 e2_encodings 保存在列表中，而这些编码仍然在 GPU 中。

解决方法：
- 在将 e1_enc 和 e2_enc 保存到列表之前，将它们转移到 CPU。
修改如下：
```
e1_encodings.append(e1_enc.cpu())
e2_encodings.append(e2_enc.cpu())
```

inputs 和 outputs 没有被释放：
尽管 torch.no_grad() 防止了计算图的构建，但 inputs 和 outputs 仍然保留在 GPU 内存中，尤其是在循环内，这些 Tensor 没有被及时释放。

解决方法：

在每次处理完后，手动删除不再使用的 Tensor，并调用 torch.cuda.empty_cache() 来释放内存。

修改如下：

e1_enc, e2_enc, tokens = get_entity_encodings(sentence)
e1_encodings.append(e1_enc.cpu())  # 将编码转移到 CPU
e2_encodings.append(e2_enc.cpu())  # 将编码转移到 CPU
sentence_labels.append(label)
sentence_indices.append(idx)

# 删除临时变量并释放显存
del e1_enc, e2_enc, tokens, inputs, outputs
torch.cuda.empty_cache()

BERT 模型的缓慢累积：
有时模型的内部可能会缓存一些中间状态，特别是当有特殊处理时（例如特殊标记的调整）。

解决方法：
- 考虑在每个批次的迭代过程中重新加载模型，以确保没有多余的缓存累积。比如，可以在每次处理一定数量的句子后，重新加载模型并继续处理。
定期监控内存并进行清理：
你可以通过手动监控 GPU 内存，检查是否有部分 Tensor 没有及时被释放。如果发现某些部分内存占用异常，可以通过分析相应的代码段来定位问题。