如何为PyTorch分配更多显存？解决RTX3090 CUDA内存不足报错

开发者阿泽

2026-7-22

解决方案

1. 调整PyTorch显存分配器参数

针对报错提示的显存碎片化问题，通过设置max_split_size_mb参数优化显存块分配逻辑：

环境变量方式（终端执行）：

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512

Python代码内设置（开头加入）：

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'

建议尝试256-1024区间的数值，找到适配你场景的最优值。

2. 主动清理显存缓存

在训练/推理的空闲节点（比如每个epoch结束后）手动释放无用资源：

import torch
import gc

# 删除不再使用的张量
del unused_tensor
# 触发Python垃圾回收
gc.collect()
# 清理PyTorch未使用的显存缓存
torch.cuda.empty_cache()

注意不要频繁调用empty_cache()，避免影响运行效率。

3. 排查隐性显存占用

确认模型、输入数据、标签都通过.to('cuda')或.cuda()正确移至GPU，避免CPU/GPU张量混合运算导致的隐性内存消耗。
推理/验证阶段必须用torch.no_grad()包裹代码，禁止不必要的梯度计算占用显存。

4. 优化训练策略

梯度累积：通过累积多步梯度再更新参数，模拟大batch效果同时降低单次显存占用：

accumulation_steps = 4
optimizer.zero_grad()
for idx, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss = loss / accumulation_steps
    loss.backward()
    if (idx + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

混合精度训练：启用自动混合精度，大幅降低显存占用：

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()
for inputs, labels in dataloader:
    optimizer.zero_grad()
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()