深度学习训练中GPU内存管理

概述

在深度学习模型训练中，主流 GPU 显存通常为 8GB~80GB，内存不足会导致训练中断或 Batch Size 受限，GPU 内存管理是优化性能和避免 OutOfMemoryError 的关键挑战。本博客简介 PyTorch 中 GPU 内存管理的核心函数、用法和实战技巧，帮助开发者高效利用显存资源。

常见问题

问题现象	解决方法
CUDA out of memory	减少 Batch Size / 使用梯度累积
显存占用持续增长	检查是否忘记 zero_grad() 或 del 变量
多卡显存不平衡	使用 torch.cuda.set_per_process_memory_fraction()
推理阶段显存过高	启用 with torch.no_grad(): 上下文管理器

1、设备选择和数据迁移

import torch

# 检查 GPU 可用性
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 将模型/数据移动到 GPU
model = model.to(device)
inputs = inputs.to(device)

# 指定特定 GPU（多卡环境）
torch.cuda.set_device(1)  # 使用编号第二块 GPU

2、显存监控函数

返回当前设备上分配的显存总量：

torch.cuda.memory_allocated(device=device)

返回在设备上分配的最大显存量：

torch.cuda.max_memory_allocated(device=device)

返回PyTorch 为了加速内存分配而预先分配的内存：

torch.cuda.memory_cached(device=device)

返回缓存显存的最大峰值

torch.cuda.max_memory_cached(device=device)

返回一个完整的内存报告，详细列出了内存的各个方面，如当前分配的内存、缓存的内存、最大显存使用情况等

torch.cuda.memory_summary(device=device)

使用示例：假设你正在训练一个大模型，你想要在每个训练周期中检查显存的分配和缓存使用情况，以避免显存溢出或缓存过度使用。你可以按如下方式进行调试：

import torch

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)

for epoch in range(num_epochs):
    print(f"Epoch {epoch}:")
    
    # 打印当前分配的显存
    print(f"  Memory Allocated: {torch.cuda.memory_allocated(device)} bytes")
    
    # 打印最大分配的显存
    print(f"  Max Memory Allocated: {torch.cuda.max_memory_allocated(device)} bytes")
    
    # 打印当前缓存的显存
    print(f"  Memory Cached: {torch.cuda.memory_cached(device)} bytes")
    
    # 打印最大缓存的显存
    print(f"  Max Memory Cached: {torch.cuda.max_memory_cached(device)} bytes")
    
    optimizer.zero_grad()

    for data, targets in dataloader:
        data, targets = data.to(device), targets.to(device)

        output = model(data)
        loss = loss_fn(output, targets)
        loss.backward()

        optimizer.step()
    
    # 打印完整的显存报告
    print(f"  Memory Summary: {torch.cuda.memory_summary(device)}")

    # 清空缓存
    torch.cuda.empty_cache()

    print("-" * 50)

3、显存释放函数

（1）torch.cuda.empty_cache()、

函数的作用是释放未使用的显存缓存。它会清理 PyTorch 的内部缓存，但不会减少 PyTorch 已经分配的显存量。

# 假设你正在训练一个大模型，可能会碰到显存不足的问题
for epoch in range(num_epochs):
    for data, targets in dataloader:
        data, targets = data.cuda(), targets.cuda()  # 将数据移动到GPU

        optimizer.zero_grad()  # 梯度清零
        output = model(data)  # 前向传播
        loss = loss_fn(output, targets)  # 计算损失
        loss.backward()  # 反向传播

        optimizer.step()  # 更新参数

        # 在训练过程中定期清理未使用的缓存
        torch.cuda.empty_cache()  # 释放空闲的缓存

（2）optimizer.zero_grad()
用来清除上一次反向传播中计算出的梯度，防止梯度的累积。如果不使用这个操作，每次 loss.backward() 后，梯度会被累加，导致梯度爆炸或计算错误。

for epoch in range(num_epochs):
    optimizer.zero_grad()  # 每个epoch开始时清除梯度
    for data, targets in dataloader:
        data, targets = data.cuda(), targets.cuda()

        # 前向传播和反向传播
        output = model(data)
        loss = loss_fn(output, targets)
        loss.backward()  # 反向传播，计算梯度

        # 执行参数更新
        optimizer.step()  # 更新模型参数

（3）del 和 torch.cuda.empty_cache()
在训练过程中，有时某些变量在计算过程中不再需要，这时可以通过删除这些变量来释放显存。del 删除的是 Python 对象，通常用于删除中间结果。结合 torch.cuda.empty_cache() 使用，可以有效地释放显存。

for epoch in range(num_epochs):
    optimizer.zero_grad()
    for data, targets in dataloader:
        data, targets = data.cuda(), targets.cuda()

        # 前向传播
        output = model(data)

        # 假设在前向传播中生成了一个临时变量
        intermediate_tensor = output.detach()

        loss = loss_fn(output, targets)  # 计算损失
        loss.backward()  # 反向传播

        optimizer.step()  # 更新参数

        # 删除不再需要的临时变量
        del intermediate_tensor  # 删除中间结果
        torch.cuda.empty_cache()  # 释放显存缓存

4、自适应batchsize调节

核心思想是从高到低检测防止内存溢出的最大batchsize值。

try:
    outputs = model(inputs)
except RuntimeError as e:  # 捕获显存不足错误
    if 'CUDA out of memory' in str(e):
        reduce_batch_size()
        retrain()

5、梯度累积

核心思想是分多次计算梯度，然后再更新模型参数，从而避免在每次迭代中都需要完整的批次数据。假设你已经有一个模型 model 和一个数据加载器 dataloader。以下是如何在 PyTorch 中实现梯度累积的示例：

import torch
from torch.optim import Adam

# 假设模型、损失函数和优化器已经定义好
model = ...
optimizer = Adam(model.parameters(), lr=1e-4)
loss_fn = ...

# 设置梯度累积的步数（例如，累积4个batch的梯度）
accumulation_steps = 4
batch_size = 32  # 这是实际使用的batch大小，可以设置较小

# 训练过程
for epoch in range(num_epochs):
    optimizer.zero_grad()  # 每个epoch开始时，先清空梯度
    for i, (data, targets) in enumerate(dataloader):
        data, targets = data.cuda(), targets.cuda()  # 假设使用GPU训练

        # 前向传播
        output = model(data)
        loss = loss_fn(output, targets)

        # 反向传播
        loss.backward()  # 计算当前batch的梯度

        # 梯度累积：如果达到了累积步数，就进行一次梯度更新
        if (i + 1) % accumulation_steps == 0:
            optimizer.step()  # 更新模型参数
            optimizer.zero_grad()  # 清空梯度

        # 如果不是累积步数的最后一步，就不更新参数，继续积累梯度