训练GitHub模型遇CUDA Out of Memory错误，常规方法无效求助

问题详情 #

runtimeerror: cuda out of memory. tried to allocate 86.00 mib (gpu 0; 4.00 gib total capacity; 3.09 gib already allocated; 0 bytes free; 3.42 gib reserved in total by pytorch)

已尝试无效方案：

• 降低训练epoch
• 执行缓存清理代码：

gc.collect()
torch.cuda.empty_cache()

设备与环境：

• 显卡：Nvidia Geforce GTX 1650（4GB显存）
• CUDA 10.2
• PyTorch 1.7.0
• Python 3.6
• 系统：Windows 10

有效解决方案 #

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

直接调整训练代码中的batch_size参数，逐步降低数值（比如从32→16→8），直到显存够用。注意：过小的batch size可能影响收敛效果，可同步按比例缩小学习率来弥补。

2. 启用混合精度训练 #

PyTorch 1.6+支持自动混合精度，通过半精度（FP16）计算大幅减少显存占用，同时几乎不损失精度。添加以下代码启用：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for epoch in range(epochs):
    for inputs, labels in dataloader:
        optimizer.zero_grad()
        # 开启混合精度上下文
        with autocast():
            outputs = model(inputs)
            loss = loss_fn(outputs, labels)
        # 缩放损失避免梯度下溢
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

3. 冻结部分模型层 #

如果是微调预训练模型，可冻结底层特征提取层，仅训练顶层任务相关层，减少需要更新的参数数量：

# 冻结所有模型参数
for param in model.parameters():
    param.requires_grad = False
# 解冻顶层需要训练的层（示例为top_layer，需替换为实际层名）
for param in model.top_layer.parameters():
    param.requires_grad = True

4. 缩小输入数据尺寸 #

针对图像类模型，降低输入图像的分辨率（比如从224×224改为192×192或128×128），直接减少单样本的显存占用。

5. 关闭PyTorch默认显存预留 #

PyTorch默认会预留部分显存，可通过代码取消限制：

torch.cuda.set_per_process_memory_fraction(1.0)  # 允许使用全部显存

Windows系统也可通过设置环境变量PYTORCH_NO_CUDA_MEMORY_CACHING=1生效。

6. 排查显存泄漏 #

检查训练代码，避免在循环内重复创建张量、模型实例，确保不再使用的变量及时被回收。比如不要在训练循环中定义模型，确保每个批次的张量都被正确处理。

内容的提问来源于stack exchange，提问作者ahmad