如何解决Ubuntu下PyTorch内存不足

1. 减少批量大小（Batch Size）

批量大小是影响GPU显存占用的核心因素之一。较小的批量大小会直接减少单次前向/反向传播所需的显存。例如，将batch_size从256降至128，显存占用可降低约50%。需注意，过小的批量可能导致训练不稳定或收敛速度变慢，需根据模型和数据集调整。

2. 使用梯度累积（Gradient Accumulation）

若无法进一步减小批量大小，可通过梯度累积模拟大批次训练。梯度累积在多个小批量上累积梯度，再执行一次参数更新，从而减少显存峰值。示例代码：

accumulation_steps = 4  # 累积4个小批次的梯度
for i, (data, target) in enumerate(dataloader):
    output = model(data)
    loss = criterion(output, target) / accumulation_steps  # 平均损失
    loss.backward()
    if (i + 1) % accumulation_steps == 0:  # 每4步更新一次参数
        optimizer.step()
        optimizer.zero_grad()

该方法可有效模拟大批次训练，同时控制显存使用。

3. 释放不必要的缓存

PyTorch会缓存未使用的显存以提高效率，但长期运行可能导致缓存堆积。通过torch.cuda.empty_cache()手动释放未使用的缓存，清理闲置显存：

import torch
torch.cuda.empty_cache()  # 清理缓存

需注意，该操作不会释放正在使用的显存，仅清理闲置部分。

4. 使用混合精度训练（Automatic Mixed Precision, AMP）

混合精度结合float16（半精度）和float32（单精度）计算，在保持模型精度的前提下，将显存占用减少约50%。PyTorch通过torch.cuda.amp模块实现自动混合精度：

from torch.cuda.amp import GradScaler, autocast
scaler = GradScaler()  # 梯度缩放器，防止数值溢出
for data, target in dataloader:
    optimizer.zero_grad()
    with autocast():  # 自动将计算转换为float16
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()  # 缩放梯度，防止溢出
    scaler.step(optimizer)  # 更新参数
    scaler.update()  # 调整缩放因子

适用于支持Tensor Cores的GPU（如NVIDIA Volta架构及以上）。

5. 优化数据加载

数据加载瓶颈会导致GPU等待，间接加剧显存压力。通过以下方式优化：

• 增加数据加载并行性：设置DataLoader的num_workers参数（如num_workers=4），利用多核CPU并行读取数据；
• 启用内存锁定：设置pin_memory=True，将数据提前锁定在物理内存中，加速传输到GPU；
• 使用生成器：对于超大规模数据集，用生成器逐条加载数据，避免一次性加载全部数据到内存。
  示例代码：

dataloader = torch.utils.data.DataLoader(
    dataset, 
    batch_size=32, 
    num_workers=4,  # 并行加载
    pin_memory=True  # 加速传输
)

6. 检查并修复内存泄漏

内存泄漏会导致显存持续增长，最终耗尽。常见原因包括：

• 循环中未释放张量（如未调用del）；
• 重复创建模型实例；
• 数据加载器未正确关闭。
  使用torch.cuda.memory_summary()监控显存变化，定位泄漏点：

print(torch.cuda.memory_summary(device=None, abbreviated=False))  # 查看显存分配详情

确保循环中及时释放不再使用的张量（如del x），并关闭数据加载器。

7. 使用更高效的模型架构

选择轻量级模型（如MobileNet、EfficientNet、SqueezeNet）替代大型模型（如ResNet-152、VGG-19），可显著减少显存占用。例如，EfficientNet-B0的参数量约为530万，而ResNet-152约为6000万，前者显存需求更低。此外，可使用模型剪枝（移除冗余神经元）、量化（将权重转换为低精度）等技术进一步压缩模型。

8. 卸载激活或参数到CPU

对于极大规模模型（如LLM），可将部分中间激活或参数暂时卸载到CPU内存，释放GPU显存。例如，在前向传播中将中间结果移至CPU，后续计算时再移回GPU：

def offload_activation(tensor):
    return tensor.cpu()  # 卸载到CPU

def process_batch(data):
    intermediate = model.layer1(data)
    intermediate = offload_activation(intermediate)  # 卸载到CPU
    intermediate = intermediate.cuda()  # 需要时移回GPU
    output = model.layer2(intermediate)
    return output

适用于无法通过上述方法解决显存瓶颈的场景。

9. 使用更精简的优化器

优化器的内存消耗差异较大。例如，Adam优化器为每个参数维护动量（m）和方差（v）两个额外状态，内存占用约为参数量的3倍；而SGD（无动量）仅维护参数本身，内存占用更小。将Adam替换为SGD，并配合余弦退火学习率调度器，可在保持收敛效果的同时减少显存：

# 替换为SGD + 余弦退火
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
num_steps = NUM_EPOCHS * len(train_loader)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_steps)

SGD的缺点是收敛速度较慢，但可通过调整学习率弥补。

10. 监控与分析显存使用

使用工具精准定位显存瓶颈：

• torch.cuda.memory_summary()：查看显存分配详情（如张量数量、大小、内存碎片）；
• nvidia-smi：实时监控GPU显存占用（如watch -n 1 nvidia-smi）；
• PyTorch Profiler：分析代码中显存消耗的热点（如某层的前向传播占用过多显存）。
  通过监控工具，可针对性优化显存占用高的部分。

声明：本文内容由网友自发贡献，本站不承担相应法律责任。对本内容有异议或投诉，请联系2913721942#qq.com核实处理，我们将尽快回复您，谢谢合作！