引言：显存监控的重要性

在深度学习模型训练过程中，显存管理直接影响训练效率和模型规模。PyTorch作为主流框架，提供了多种显存监控工具，但开发者常因不了解具体方法而陷入显存泄漏或不足的困境。本文将系统梳理PyTorch中显存查看与监控的核心技术，从基础API到高级工具，为不同场景提供解决方案。

一、基础显存查看方法

1.1 使用torch.cuda模块

PyTorch通过torch.cuda子模块提供显存信息查询功能，核心方法包括：

import torch
# 查看当前GPU显存总量（MB）
total_memory = torch.cuda.get_device_properties(0).total_memory / 1024**2
print(f"Total GPU Memory: {total_memory:.2f} MB")
# 查看当前显存占用（MB）
allocated_memory = torch.cuda.memory_allocated() / 1024**2
reserved_memory = torch.cuda.memory_reserved() / 1024**2
print(f"Allocated: {allocated_memory:.2f} MB, Reserved: {reserved_memory:.2f} MB")

• 关键指标解析：
  • memory_allocated()：当前被PyTorch张量占用的显存
  • memory_reserved()：CUDA缓存分配器保留的显存（包含未使用部分）
  • total_memory：GPU物理显存总量

1.2 显存使用峰值追踪

通过torch.cuda.max_memory_allocated()可获取训练过程中的显存峰值：

# 在训练循环前后调用
before_max = torch.cuda.max_memory_allocated() / 1024**2
# ...执行训练步骤...
after_max = torch.cuda.max_memory_allocated() / 1024**2
print(f"Memory peak increased by: {after_max - before_max:.2f} MB")

此方法特别适用于定位显存泄漏点，通过比较不同训练阶段的峰值变化，可快速定位异常内存增长。

二、动态显存监控技术

2.1 训练过程中的实时监控

结合torch.cuda与自定义回调函数，可实现训练过程中的动态监控：

class MemoryMonitor:
    def __init__(self):
        self.history = []
    def __call__(self, engine):
        mem = torch.cuda.memory_allocated() / 1024**2
        self.history.append(mem)
        print(f"Step {engine.state.epoch}: {mem:.2f} MB")
# 在ignite训练引擎中使用
from ignite.engine import Engine
monitor = MemoryMonitor()
engine = Engine(train_step)
engine.add_event_handler("iteration_completed", monitor)

此方案适用于需要精确追踪每步显存变化的场景，如调试复杂模型结构时的内存消耗模式。

2.2 使用NVIDIA工具扩展监控

NVIDIA提供的nvtop和nvidia-smi命令行工具可与PyTorch监控形成互补：

# 实时监控显存使用（需安装nvidia-tools）
nvidia-smi -l 1 --query-gpu=memory.used,memory.free --format=csv

• 优势对比：
  • PyTorch API：精确到张量级别的内存分配
  • NVIDIA工具：显示所有进程的显存占用，包括非PyTorch进程

三、高级显存分析技术

3.1 内存分析器（Memory Profiler）

PyTorch 1.10+版本内置的内存分析器可生成详细内存使用报告：

with torch.autograd.profiler.profile(
    use_cuda=True,
    profile_memory=True,
    record_shapes=True
) as prof:
    # 执行需要分析的代码
    output = model(input_tensor)
print(prof.key_averages().table(
    sort_by="cuda_memory_usage", 
    row_limit=10
))

输出结果包含每个操作的显存分配量、操作类型和调用栈，是定位显存瓶颈的利器。

3.2 可视化监控工具

对于复杂项目，推荐使用以下可视化方案：

1. TensorBoard集成：
   ```python
   from torch.utils.tensorboard import SummaryWriter
   writer = SummaryWriter()

def log_memory(step):
mem = torch.cuda.memory_allocated() / 1024**2
writer.add_scalar(“Memory/Allocated”, mem, step)

2. **PyTorch Profiler的Chrome追踪**：
```python
prof = torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CUDA],
    profile_memory=True
)
with prof:
    # 执行代码
prof.export_chrome_trace("trace.json")

通过Chrome的chrome://tracing加载生成的JSON文件，可获得时间轴上的显存分配可视化。

四、显存优化实践建议

4.1 常见问题诊断

• 显存泄漏：特征为memory_allocated()持续上升
  • 解决方案：检查循环中的张量是否被正确释放，使用del tensor和torch.cuda.empty_cache()
• 碎片化问题：表现为reserved_memory远大于allocated_memory
  • 解决方案：启用torch.backends.cuda.cufft_plan_cache.clear()清理缓存

4.2 最佳实践

1. 混合精度训练：

   scaler = torch.cuda.amp.GradScaler()
   with torch.cuda.amp.autocast():
    outputs = model(inputs)

   可减少显存占用达50%

2. 梯度检查点：
   ```python
   from torch.utils.checkpoint import checkpoint
   def custom_forward(*inputs):

   实现前向传播

   return outputs

outputs = checkpoint(custom_forward, *inputs)

适用于深层网络，以计算换内存
3. **模型并行**：
```python
# 将模型分到不同GPU
model = nn.DataParallel(model, device_ids=[0,1,2,3])

五、跨平台兼容性考虑

5.1 多GPU环境监控

在DDP（Distributed Data Parallel）环境下，需指定设备ID：

def get_gpu_memory(device_id=0):
    allocated = torch.cuda.memory_allocated(device_id) / 1024**2
    reserved = torch.cuda.memory_reserved(device_id) / 1024**2
    return allocated, reserved

5.2 CPU与GPU显存对比

对于CPU上的内存监控，可使用resource模块：

import resource
def get_cpu_memory():
    return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024  # MB

结论与展望

显存监控是深度学习工程化的核心能力，本文系统介绍了从基础API到高级分析工具的完整方案。实际应用中，建议采用分层监控策略：

1. 开发阶段使用torch.cudaAPI进行快速检查
2. 调试阶段结合内存分析器和可视化工具
3. 生产环境部署TensorBoard或Prometheus监控

未来随着PyTorch生态的发展，预计会出现更智能的显存管理方案，如自动碎片整理和动态批处理优化。开发者应持续关注框架更新，保持显存管理技术的先进性。