【Pytorch使用】CUDA 显存管理与 OOM 排查实战：以 PyTorch 联邦学习训练为例

Rabbit_QL2026-01-07 13:19

一、问题背景

​ 在深度学习训练过程中，特别是使用 GPU 进行联邦学习训练时，经常会遇到 torch.OutOfMemoryError: CUDA out of memory 错误。

1.1 现象发生的可能原因

  • 模型较大（如 VGG16、ResNet 等）
  • 批次大小（batch_size）过大
  • 多个客户端共享模型对象
  • 长时间训练导致内存碎片化
  • 未及时释放中间变量和缓存

1.2 典型错误信息

bash 复制代码 
torch.OutOfMemoryError: CUDA out of memory. 
Tried to allocate 6.12 GiB. 
GPU 2 has a total capacity of 44.42 GiB of which 1.67 GiB is free. 
Including non-PyTorch memory, this process has 42.74 GiB memory in use. 
Of the allocated memory 21.18 GiB is allocated by PyTorch, 
and 21.23 GiB is reserved by PyTorch but unallocated.

​ GPU 显存不足：当前进程已经几乎占满 GPU 显存，PyTorch 尝试再分配一块 6.12 GiB 的显存时失败。

项目 显存
GPU 总显存 44.42 GiB
当前可用 1.67 GiB
进程占用总量 42.74 GiB
PyTorch 实际使用 21.18 GiB
PyTorch 缓存预留 21.23 GiB
本次尝试申请 6.12 GiB ❌
  • 该 OOM 错误表明，尽管 GPU 总显存仍有少量空闲，但由于 PyTorch 显存缓存机制与当前显存碎片化，无法再为新的计算分配一块连续的 6.12 GiB 显存，从而触发 CUDA out of memory 错误。

二、CUDA 内存不足的原因分析

2.1 内存分配问题

  • 问题描述：

    • PyTorch 使用内存池（memory pool）管理 GPU 内存
    • 当内存碎片化严重时，即使有足够的空闲内存，也可能无法分配连续的大块内存

  • 原因：

    • 频繁的小块内存分配和释放
    • 未及时清理缓存
    • 多个模型实例同时占用内存

2.2 批次大小过大

  • 问题描述：

    • batch_size 过大导致单次前向传播需要大量显存
    • VGG16 模型在 224 × 224 224 \times 224 224×224 输入下，每个样本约需 12 − 15 M B 12-15 MB 12−15MB 显存
    • batch_size=512 需要约 6 − 7.5 G B 6-7.5 GB 6−7.5GB 显存（仅前向传播）

  • 计算公式：
    KaTeX parse error: Expected 'EOF', got '_' at position 30: ...} = \text{batch_̲size} \times \t...

2.3 模型对象共享

问题描述：

  • 多个客户端共享同一个模型对象引用
  • 在联邦学习中，所有客户端同时训练会导致内存累积
  • 模型参数的梯度信息会占用额外内存

2.4 计算图保留

问题描述：

  • 默认情况下，PyTorch 会保留计算图用于反向传播
  • 在计算参数差异时，如果不需要梯度，应该使用 torch.no_grad()

三、解决方案概述

​ 我们的解决方案采用多层次的内存优化策略：

  • 环境变量优化 - 减少内存碎片

  • 批次大小调整 - 降低单次内存需求

  • 模型副本管理 - 避免共享对象导致的内存累积

  • 主动内存清理 - 及时释放不需要的内存

  • 计算图优化 - 避免不必要的梯度计算

3.1 环境变量优化

  • 方法：

    在脚本开头设置环境变量，优化 CUDA 内存分配策略：

    python 复制代码 
    import os
# 优化 CUDA 内存分配，避免内存碎片
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

  • 作用：

    • expandable_segments:True 允许内存段动态扩展，减少内存碎片
    • 提高内存利用率，避免"有内存但无法分配"的情况

  • 参考文档： PyTorch Memory Management

3.2 批次大小调整

​ 根据模型大小和 GPU 显存容量选择合适的 batch_size：

python 复制代码 
conf = {
    "batch_size": 32,  
    # VGG16 + 224x224: 推荐 16-64
    # ResNet50 + 224x224: 推荐 32-128
    # 小型模型: 可以更大
}
(1) 选择原则
  • 大型模型（VGG16, ResNet50）: 16-64
  • 中型模型（ResNet18）: 32-128
  • 小型模型（MobileNet）: 64-256
  • 输入尺寸: 224x224 使用较小值，32x32 可以使用较大值
(2)内存估算
python 复制代码 
# 粗略估算公式
estimated_memory = batch_size * input_size * channels * 4 bytes * 3
# 3 倍包括：输入、梯度、优化器状态

3.3 模型副本管理

​ 为每个客户端创建独立的模型实例，避免共享对象：

python 复制代码 
from torchvision import models
import torch.nn as nn

# 为每个客户端创建独立的模型副本
client_model = models.vgg16(weights=None)
# 修改分类层
in_features = client_model.classifier[6].in_features
client_model.classifier[6] = nn.Linear(in_features, num_classes)
# 复制参数（不共享对象）
client_model.load_state_dict(server.global_model.state_dict())
client_model = client_model.to(device)
(1)关键点
  • ✅ 正确做法: 创建新模型实例，只复制参数
  • ❌ 错误做法 : 使用 copy.deepcopy() 或直接共享模型对象
(2)内存对比
  • 共享模型对象：所有客户端共享同一份模型内存 + 梯度内存累积
  • 独立模型副本：每个客户端独立内存，训练完成后可释放

3.4 主动内存清理

​ 在关键位置主动清理 GPU 缓存，但是需要注意频繁调用 empty_cache() 会引入同步开销，不建议在每个 batch 调用。

(1) 训练循环中的清理
python 复制代码 
for batch_id, (imgs, labels) in enumerate(self.train_loader):
    data, target = imgs.to(self.device), labels.to(self.device)
    optimizer.zero_grad()
    output = self.local_model(data)
    loss = loss_func(output, target)
    loss.backward()
    optimizer.step()
    
    # 清理中间变量释放内存
    del data, target, output, loss
    
    # 定期清理缓存
    if (batch_id + 1) % 10 == 0:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
  • 清理时机：
    • 每个批次后删除中间变量
    • 每 N 个批次清理一次缓存（N=10 或根据内存情况调整）
    • 客户端训练完成后清理
(2) 参数差异计算优化
python 复制代码 
# 使用 no_grad 避免保存计算图
with torch.no_grad():
    diff = dict()
    for name, data in self.local_model.state_dict().items():
        global_data = global_model.state_dict()[name]
        # 移到 CPU 避免 GPU 内存累积
        diff[name] = (data - global_data).cpu()

# 清理 GPU 缓存
if torch.cuda.is_available():
    torch.cuda.empty_cache()
  • 优化点：
    • torch.no_grad(): 不保存计算图，节省内存
    • .cpu(): 将结果移到 CPU，释放 GPU 内存
    • 计算完成后立即清理缓存
(2) 模型聚合优化
python 复制代码 
def model_aggregate(self, weight_accumulator):
    with torch.no_grad():
        for name, data in self.global_model.state_dict().items():
            # 确保张量在正确的设备上
            if isinstance(weight_accumulator[name], torch.Tensor):
                update_per_layer = weight_accumulator[name].to(data.device) * self.conf["lambda"]
            else:
                update_per_layer = weight_accumulator[name] * self.conf["lambda"]
            data.add_(update_per_layer)
    
    # 清理 GPU 缓存
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

3.5 内存监控

​ 添加内存监控代码，便于调试和优化：

python 复制代码 
if torch.cuda.is_available():
    device_id = 2
    torch.cuda.set_device(device_id)
    device = torch.device(f'cuda:{device_id}')
    
    # 清理 GPU 缓存
    torch.cuda.empty_cache()
    
    # 显示内存使用情况
    allocated = torch.cuda.memory_allocated(device_id) / 1024**3
    total = torch.cuda.get_device_properties(device_id).total_memory / 1024**3
    print(f"GPU 内存使用情况: {allocated:.2f} GB / {total:.2f} GB")

四、最佳实践

4.1 批次大小选择

python 复制代码 
# 根据 GPU 显存选择 batch_size
gpu_memory_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3

if gpu_memory_gb >= 24:  # RTX 3090, A100 等
    batch_size = 64
elif gpu_memory_gb >= 12:  # RTX 3060, V100 等
    batch_size = 32
else:  # GTX 1080 等
    batch_size = 16

4.2 梯度累积（Gradient Accumulation）

​ 当显存不足但需要大批次训练时，使用梯度累积：

python 复制代码 
accumulation_steps = 4  # 累积 4 个批次
effective_batch_size = batch_size * accumulation_steps

for batch_id, (data, target) in enumerate(train_loader):
    output = model(data)
    loss = criterion(output, target) / accumulation_steps
    loss.backward()
    
    if (batch_id + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()
        torch.cuda.empty_cache()

4.3 混合精度训练（Mixed Precision）

使用 FP16 可以显著减少显存占用：

python 复制代码 
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for data, target in train_loader:
    optimizer.zero_grad()
    
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

4.4 检查点机制（Checkpointing）

对于非常大的模型，可以使用梯度检查点：

python 复制代码 
from torch.utils.checkpoint import checkpoint

# 在前向传播中使用检查点
output = checkpoint(model, input)

4.5 内存清理检查清单

在以下位置添加内存清理：

✅ 每个训练批次后删除中间变量

✅ 每个客户端训练完成后清理缓存

✅ 模型聚合后清理缓存

✅ 评估前清理缓存

✅ 使用 torch.no_grad() 进行不需要梯度的计算

五、故障排查

问题 1: 仍然出现内存不足

可能原因：

  • batch_size 仍然太大
  • 模型太大
  • 多个进程共享 GPU

解决方案：

python 复制代码 
# 1. 进一步减小 batch_size
conf["batch_size"] = 16  # 或更小

# 2. 使用梯度累积
# 3. 使用混合精度训练
# 4. 检查是否有其他进程占用 GPU
nvidia-smi  # 查看 GPU 使用情况

问题 2: 内存清理后仍然占用大量内存

可能原因：

  • PyTorch 的内存池机制会保留已分配的内存
  • 这是正常现象，不影响使用

解决方案：

python 复制代码 
# 如果需要强制释放，可以重启 Python 进程
# 或者在代码中定期重置
torch.cuda.empty_cache()
torch.cuda.ipc_collect()  # 收集 IPC 资源

问题 3: 训练速度变慢

可能原因：

  • 频繁的内存清理会影响性能
  • batch_size 太小

解决方案：

python 复制代码 
# 1. 减少清理频率
if (batch_id + 1) % 50 == 0:  # 改为每 50 个批次清理一次
    torch.cuda.empty_cache()

# 2. 在非关键路径清理
# 3. 使用更大的 batch_size（如果内存允许）

问题 4: 多 GPU 训练时的内存问题

解决方案：

python 复制代码 
# 使用 DataParallel 或 DistributedDataParallel
if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)
    # 注意：batch_size 会按 GPU 数量自动分配

六、代码示例总结

6.1 完整的客户端训练代码示例

python 复制代码 
def local_train(self, global_model):
    # 1. 复制模型参数（使用 no_grad）
    with torch.no_grad():
        self.local_model.load_state_dict(global_model.state_dict())
    
    # 2. 初始化优化器
    optimizer = torch.optim.SGD(self.local_model.parameters(), lr=0.001)
    loss_func = torch.nn.CrossEntropyLoss()
    
    # 3. 训练循环
    self.local_model.train()
    for epoch in range(self.conf["local_epochs"]):
        for batch_id, (imgs, labels) in enumerate(self.train_loader):
            # 数据移动到设备
            data, target = imgs.to(self.device), labels.to(self.device)
            
            # 前向传播和反向传播
            optimizer.zero_grad()
            output = self.local_model(data)
            loss = loss_func(output, target)
            loss.backward()
            optimizer.step()
            
            # 清理中间变量
            del data, target, output, loss
            
            # 定期清理缓存
            if (batch_id + 1) % 10 == 0:
                torch.cuda.empty_cache()
    
    # 4. 计算参数差异（移到 CPU）
    with torch.no_grad():
        diff = dict()
        for name, data in self.local_model.state_dict().items():
            diff[name] = (data - global_model.state_dict()[name]).cpu()
    
    # 5. 最终清理
    torch.cuda.empty_cache()
    return diff

参考资料

  1. PyTorch CUDA Memory Management
  2. PyTorch Performance Tuning Guide
  3. Understanding CUDA Memory Allocation
相关推荐
一水鉴天2 小时前
映射、哈希表与哈斯图：计算机科学的三种基线 20261003（元宝）
开发语言·人工智能
198******126342 小时前
2026 企业 AI 办公产品选型指南：从场景匹配判断工具价值
人工智能
玫瑰互动GEO3 小时前
GEO优化学习九级模型：开发者从认知层切入
人工智能·ai·ai搜索·gem·生成式引擎优化·gem优化
海绵宝宝转agent3 小时前
learn-claude-code第1-5章开源学习笔记分享
人工智能·笔记·python·学习
小奇不哭3 小时前
ROS2 单节点，RGB-D 深度相机局部路径提取 + B 样条平滑路径 + YOLOv8 障碍物检测 + 深度测距，用于移动机器人视觉局部循迹。
python·yolov8·ros2·cv2·路径探索循迹
每天都要写算法（努力版）3 小时前
【行业前沿报告】DAgger：让智能体在自己走到的状态上学习
人工智能·学习·机器学习
liron713 小时前
智能实体演化系统的统一性概念
人工智能·深度学习·神经网络
小小龙学IT3 小时前
Python scikit-learn 机器学习库深度解析
python·机器学习·scikit-learn
呆萌很3 小时前
常用骨干网络预训练输入尺寸
人工智能
Yolanda_20223 小时前
19.神经网络-最大池化的使用
人工智能·深度学习·神经网络
热门推荐
01国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）02历年考研数学一、数学二、数学三真题试卷及答案PDF03如何新建文件夹？ 电脑新建文件夹的4种方法04Telegram APK 下载安装完整指南 — 2026年最新05GitHub 镜像站点06Windows 10/11 如何退出微软账户并改用本地账户登录（含找不到选项的解决方法）07蜘蛛磁力 搜索引擎大全，如何使用蜘蛛磁力查找磁力链接08微信历史版本含下载地址（ Windows PC | 安卓 | MAC ）及设置微信不更新092026 爱知・名古屋亚运会怎么看？直播平台+赛程安排分享10给AI一张“业务概念地图“：本体如何从哲学走向企业智能