训练SciBERT验证阶段CUDA显存不足问题的解决方法咨询

我用约10000行的数据集训练SciBERT，按0.6、0.2、0.2的比例划分训练集、验证集和测试集，训练代码如下：

from torch.optim import AdamW
from tqdm import tqdm
from statistics import mean

optim = AdamW(model.parameters(), lr=2e-5, eps=1e-8)

for epoch in range(4):
    epoch_losses = []
    validation_loss = []
    for x, y in tqdm(load_data(x_train, y_train, batch_size=10)):
        model.zero_grad()
        out = model(x, attention_mask=apply_attention_mask(x), labels=y)
        epoch_losses.append(out.loss.item())
        out.loss.backward()
        optim.step()
    print(f"epoch {epoch + 1} loss: {mean(epoch_losses)}")

    for x, y in load_data(x_validation, y_validation, batch_size=10):
    # validation data
        validation_output = model(x, attention_mask=apply_attention_mask(x), labels=y)
        validation_loss.append(validation_output.loss.item())
    print(f"Validation for epoch {epoch + 1} loss: {mean(validation_loss)}")

但完成一轮训练集训练后加载验证集时，Colab返回如下错误：

OutOfMemoryError: CUDA out of memory. Tried to allocate 120.00 MiB (GPU 0; 14.75 GiB total capacity; 13.84 GiB already allocated; 24.81 MiB free; 14.59 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

以下是具体修复方案：

核心修复步骤 #

• 验证阶段关闭梯度计算：训练阶段产生的梯度会持续占用显存，验证时不需要反向传播，用torch.no_grad()上下文管理器包裹验证循环，彻底避免梯度相关的显存消耗。修改后的验证代码片段：

  with torch.no_grad():
      for x, y in load_data(x_validation, y_validation, batch_size=10):
          # validation data
          validation_output = model(x, attention_mask=apply_attention_mask(x), labels=y)
          validation_loss.append(validation_output.loss.item())
  

• 手动清理显存缓存：在训练循环结束后、验证开始前，清理未使用的张量和显存碎片，释放被占用的冗余显存。添加代码：

  import torch
  torch.cuda.empty_cache()
  

  放在训练循环结束后，验证循环之前。

• 调整验证批次大小：如果显存还是紧张，可适当降低验证集的batch_size（比如从10改为8或5），减少单步推理的显存占用。

• 确保模型仅加载一次：确认训练前已将模型移至GPU，避免验证阶段重复加载。在代码开头添加：

  model = model.to('cuda')
  

整合优化后的完整代码 #

import torch
from torch.optim import AdamW
from tqdm import tqdm
from statistics import mean

# 确保模型加载到GPU
model = model.to('cuda')
optim = AdamW(model.parameters(), lr=2e-5, eps=1e-8)

for epoch in range(4):
    epoch_losses = []
    validation_loss = []
    
    # 训练阶段
    model.train()
    for x, y in tqdm(load_data(x_train, y_train, batch_size=10)):
        model.zero_grad()
        out = model(x, attention_mask=apply_attention_mask(x), labels=y)
        epoch_losses.append(out.loss.item())
        out.loss.backward()
        optim.step()
    print(f"epoch {epoch + 1} loss: {mean(epoch_losses)}")
    
    # 清理显存缓存
    torch.cuda.empty_cache()
    
    # 验证阶段
    model.eval()
    with torch.no_grad():
        for x, y in load_data(x_validation, y_validation, batch_size=10):
            validation_output = model(x, attention_mask=apply_attention_mask(x), labels=y)
            validation_loss.append(validation_output.loss.item())
    print(f"Validation for epoch {epoch + 1} loss: {mean(validation_loss)}")

内容的提问来源于stack exchange，提问作者Hoang Cuong Nguyen