1. Pytorch显存分配机制概述
Pytorch是流行的深度学习框架之一，它使用动态计算图系统，可以通过计算图来追踪神经网络的计算流程。然而，在使用Pytorch训练深度神经网络时，由于深度神经网络存在大量的参数和中间计算结果，因此很容易导致显存不足的问题。因此，了解Pytorch显存的分配机制非常重要。
2. Pytorch显存管理器
2.1 CUDA管理中的内存分配
在Pytorch中，CUDA提供了内存管理功能，Pytorch显存管理器可以使用这些功能来管理显存的分配和释放。每个Pytorch张量都有一个指针指向其在显存中的地址。当Pytorch需要为一个张量分配显存时，它将调用CUDA库中的分配器函数来分配一段内存，然后将张量的指针设置为该内存的地址。同样地，当张量不再需要时，它的内存将由CUDA分配器函数释放。
import torch
a = torch.randn(5000,5000, device='cuda:0')
print(torch.cuda.memory_allocated())
其中的memory_allocated()函数用来检查当前显存的分配情况。
2.2 张量的共享内存
在Pytorch中，多个张量可以共享同一块显存。这种情况通常会出现在深度神经网络的反向传播过程中，因为反向传播需要反向遍历整个计算图，其中一些中间结果可能会被多个张量引用。
2.3 Pytorch的存储管理
在Pytorch中，存储被用来存储张量的数据。张量的数据分为两部分，一部分是用来存储实际数据的内存块，另一部分是用来跟踪张量形状和步幅的元数据。
3. Pytorch显存的分配策略
3.1 延迟分配
在Pytorch中，默认情况下，张量的显存并不会立即分配。Pytorch使用了延迟分配的策略，只有当首次访问张量的数据时，Pytorch才会为其分配显存。这个机制可以有效地减少显存占用，特别是在处理大型模型时。
3.2 梯度缓存的释放
在训练深度神经网络时，Pytorch中的自动梯度计算机制会在反向传播时缓存梯度，在反向遍历计算图结束后，Pytorch会清空梯度缓存，释放相关的显存。
3.3 内存碎片的问题
由于动态计算图的机制，Pytorch能够支持任意形状和大小的张量，但这也导致显存碎片的问题。尤其是在模型权重变小、变大或者形状改变时，Pytorch的显存管理器需要将内存块重新整理，以便生成连续的内存块。这个过程可能会导致不必要的显存占用。
4. Pytorch显存管理技巧
4.1 在GPU和CPU之间移动数据
在训练深度神经网络时，需要将数据送入模型进行计算，同时需要将模型的输出送回至主机。在Pytorch中，可以使用to()方法来在GPU和CPU之间移动数据。
import torch
a = torch.randn(5000,5000)
a_cpu = a.to('cpu')
a_gpu = a.to('cuda:0')
4.2 清空缓存
在训练模型时，Pytorch会缓存许多中间值，这些缓存会占用显存。可以使用torch.cuda.empty_cache()方法来清空Pytorch缓存。
import torch
a = torch.randn(5000,5000, device='cuda:0')
print(torch.cuda.memory_allocated())
torch.cuda.empty_cache()
print(torch.cuda.memory_allocated())
4.3 减少batch size
可以通过减少batch size的大小来减少显存占用。然而，在减小batch size的同时应该注意训练速度的下降。
5. 总结
本文介绍了Pytorch显存管理器的工作原理、分配策略以及一些管理技巧。理解这些内容可以帮助我们更好地利用显存资源，避免因显存不足而影响模型训练的效率。