PyTorch显存机制简要分析

1. 显存开销：显存占用约为element个数的4倍，注意不会将空闲的显存释放
   • 模型参数：
   • 模型参数的梯度
   • 优化器状态：跟踪每个权重参数，比如记录动量
   • 中间结果：前向传播时计算的中间结果

2. 显存分析方法
   1. No Nvidia-smi：PyTorch中即使一个Tensor被释放了，进程也不会将空闲出的显存释放，而是等待新的Tensor分配，这样进程不需要重新向GPU申请显存。因此使用Nvidia-smi得到的显存是reserved memory（已经分配到的缓冲区显存，但是并不一定正在使用）+ PyTorch context（PyTorch在创建进程cuda时所需的开销）
   2. reversed memory
      • PyTorch创建的torch.Tensor都是在global memory上，但是使用一些方法（比如cublas）可以使用shared memory等进行加速
      • Pytorch显存管理机制
      • use torch.cuda
        • torch.cuda.memory_allocated()：当前进程torch.Tensor占用的显存
        • torch.cuda.max_memory_allocated()：到调用函数为止所达到的最大显存占用，可以搭配reset_peak_memory_stats()函数使用
        • torch.cuda.memory_reserved：caching allocator管理的显存
        • torch.cuda.max_memory_reserved：到调用函数为止caching allocator管理的显存
        • torch.cuda.memory_summary
        • torch.cuda.empty_cache()：释放caching allocator占用的空闲（unused）显存
      • useful tools: pytorch_memlab Profiling and inspecting memory in pytorch

3. PyTorch context
   • 比如在cuda上创建了一个很小的数组，但是nvidia-smi占用了很多显存（大约1000M左右）

4. 显存分配机制
   1. 显存按页分配，PyTorch先向CUDA device申请2MB的显存放到自己的cache（使用torch.cuda.memory_reserved查看）
   2. 创建torch.Tensor时从cache中分配空间（使用torch.cuda.memory_allocated查看）

5. 显存释放机制
   1. 在PyTorch中，只要一个Tensor对象在后续不会再被使用，那么PyTorch就会自动回收该Tensor所占用的显存，并以缓冲区的形式继续占用显存。

6. 参考

   主要参考：PyTorch显存机制分析

   一文读懂PyTorch显存管理机制

   A comprehensive guide to memory usage in PyTorch

   Deep leaning memory usage and pytorch optimization tricks

   How to locate data location: global memory or shared memory?