问题：PyTorch中如何实时查看GPU显存使用情况？

Question

在使用PyTorch进行深度学习训练或推理时，如何实时查看GPU显存使用情况是一个常见且关键的问题。许多开发者在调试模型、优化内存占用或排查显存溢出（Out of Memory）错误时，迫切需要了解当前GPU显存的分配和使用状况。PyTorch提供了多种内置方法和工具，如`torch.cuda.memory_allocated()`、`torch.cuda.memory_reserved()`等函数，可用于程序中实时监控显存使用。此外，结合第三方工具如NVIDIA的`nvidia-smi`命令行工具或PyTorch Profiler，也能实现更直观和全面的监控。掌握这些方法对于提升模型训练效率和资源管理至关重要。

Answer 1

一、PyTorch中GPU显存监控的初步认识

在深度学习训练或推理过程中，GPU显存的使用情况直接影响模型的运行效率与稳定性。PyTorch作为主流深度学习框架之一，提供了丰富的API用于监控显存使用。

常见的显存监控函数包括：

• torch.cuda.memory_allocated()：返回当前分配的显存大小（单位为字节）。
• torch.cuda.memory_reserved()：返回当前保留的显存大小（包括已分配和未分配部分）。
• torch.cuda.max_memory_allocated()：返回程序运行过程中分配过的最大显存。
• torch.cuda.empty_cache()：释放未被占用的缓存，有助于调试显存问题。

二、显存监控的进阶用法

为了更深入地理解显存使用情况，开发者可以将这些函数嵌入到训练或推理的每一步中，进行实时监控。例如，在训练循环中打印显存信息：

for epoch in range(epochs):
    for inputs, labels in dataloader:
        inputs, labels = inputs.cuda(), labels.cuda()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        print(f"Allocated: {torch.cuda.memory_allocated() / 1024 ** 2:.2f} MB")
        print(f"Reserved: {torch.cuda.memory_reserved() / 1024 ** 2:.2f} MB")
    

通过这种方式，开发者可以发现显存使用的峰值，并据此优化模型结构或批量大小。

三、结合第三方工具进行系统级监控

除了在代码中进行显存监控外，还可以使用系统级工具进行更全面的观察。

NVIDIA的nvidia-smi命令行工具可以实时查看GPU的显存使用、温度、功耗等信息：

$ nvidia-smi -q -d MEMORY
    

输出示例（简化）：

GPU ID	Used Memory	Total Memory
0	3072 MB	16384 MB
1	2048 MB	16384 MB

四、使用PyTorch Profiler进行细粒度分析

PyTorch Profiler是PyTorch提供的性能分析工具，可以记录每一步操作的显存使用情况，帮助开发者识别内存瓶颈。

with torch.profiler.profile(profile_memory=True, record_shapes=True) as prof:
    output = model(input)
print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))
    

输出示例（简化）：

Operator	CUDA Memory Usage (MB)
aten::linear	150.2
aten::addmm	98.5
aten::relu	20.1

五、显存监控的可视化与自动化

对于需要长期运行或自动监控的场景，可以将显存数据写入日志文件或使用可视化工具进行展示。

例如，使用TensorBoard记录显存使用情况：

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()

for step in range(100):
    writer.add_scalar('GPU Memory Allocated', torch.cuda.memory_allocated() / 1024**2, step)
    writer.add_scalar('GPU Memory Reserved', torch.cuda.memory_reserved() / 1024**2, step)
    # 模拟训练过程
    time.sleep(0.1)
    

通过TensorBoard，开发者可以直观地观察显存使用趋势，辅助优化模型训练过程。

六、显存溢出（OOM）的排查与优化策略

当出现显存溢出错误时，以下策略可以帮助排查与优化：

1. 减少批量大小（batch size）。
2. 使用混合精度训练（AMP）。
3. 启用torch.utils.checkpoint以空间换时间。
4. 使用torch.cuda.empty_cache()释放未使用的缓存。
5. 检查模型中是否存在冗余计算或内存泄漏。

此外，结合PyTorch Profiler和nvidia-smi可以更准确地定位问题源头。

七、监控流程图示例

下图展示了在PyTorch训练流程中嵌入显存监控的逻辑流程：

graph TD
    A[开始训练] --> B[加载数据]
    B --> C[前向传播]
    C --> D[反向传播]
    D --> E[更新参数]
    E --> F[调用显存监控函数]
    F --> G{是否完成训练?}
    G -- 否 --> A
    G -- 是 --> H[结束训练]