|
今天我们来谈谈使用GPU时，常常会面临的一个内存不足的问题，以及如何解决。
当我们在训练较大的深度学习模型时，你很快就会发现，你花了那么多钱买的炫酷的GPU（或者可能更明智地使用了云实例上的GPU）会经常出问题，总是抱怨内存不足。不过，要知道GPU有数GB的内存！怎么可能不足呢？
模型往往会占用大量内存。例如，ResNet-152有约6000万激活，所有这些都会占用GPU上宝贵的空间。下面我们来看看GPU内部，确定内存不足时可能在做什么。
检查GPU
如果你使用了一个NVIDIA GPU（如果你使用的是其他显卡，需要查看该GPU供应商的驱动程序网站来得到相应工具），CUDA安装包含有一个很有用的命令行工具，名为nvidia-smi。运行时如果没有提供参数，这个工具会给出GPU上所用内存的一个快照，更好的一点是，还会给出谁在使用这些内存。
如果确定认为你的GPU占用的内存超过了应有的使用量，可以尝试让Python的垃圾回收器介入。如果有一个不再需要的tensor_to__be_deleted,你想将它从GPU删除，fast.ai库提供的方法是用del删除：
import gc
del tensor_to_be_deleted
gc.collect()
梯度检查点
尽管可以使用上一节的删除和垃圾回收技巧，但你可能还是发现内存不足。对于大多数应用来说，但是这样一来，就会增加每个epoch的训练时间，而且与有足够内存来处理更大批量数据的模型相比，这个模型可能稍逊一筹，因为，前者每次传播可以看到数据集中的更多数据。不过，在PyTorch中可以使用梯度检查点(gradient checkpointing),用计算量为大模型换取内存。
处理更大的模型时，问题之一就是前向和反向传播会创建大量中间状态，所有这些都会占用GPU内存。内存检查点的目的就是通过分段（segmenting）减少可能存放在GPU上的状态数量。这个方法意味着，模型的批量大小可以达到非分段模型的4到10倍，但与此同时，训练的计算量更大。在前向传播中，PyTorch将输入和参数存储到一个分段，但本身并不具体完成前向传播。在反向传播中，由PyTorch获取这些输入和参数，为该分段计算前向传播。中间值传递到下一个分段，但这些只能一个分段一个分段地完成。
将一个模型切分为这些分段是由torch.utils.checkpoint.checkpoint_sequential()处理的。它作用于nn.Sequential层或生成的层列表，条件是需要按照它们在模型中出现的顺序进行处理。
为了更清楚的说明，这里放一段代码进行说明：
from torch.utils.checkpoint import chechpoint_sequential
sefl.avgpool=nn.AdaptiveAvgPool2d((6,6))
self.classifier=nn.Sequential(
nn.Dropout(),
nn.Linear(256*6*6,4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096,4096),
nn.ReLU(inplace=True),
nn.Linear(4096，num_classes),
)
def forward(self,x):
x=checkpoint_sequential(self.features,chunks,x)
x=self.avgpool(x)
x=x.view(x.size(0),256*6*6)
x=self.classifier(x)
return x
可以看到，这里并没有太多不同，所以必要时，可以很容易地为模型增加检查点。我们为这个新版本的模型增加一个chunks参数，默认地将它分为两个分段。然后所要作的就是调用checkpoint_sequential并传入features模块，分段数以及我们的输入。
检查点技术有一个小问题：对于BatchNorm或者Dropout层，它的表现不太好，原因在于这些层于前向传播的交互方式。为了避免这个问题，可以对这些层之前或者之后的模型部分增加检查点。在我们的ChenckpointedAlexNet中，可以把classifier模块分解为两部分：包含Dropout层的部分不加检查点，另一个部分是最后的一个nn.Sequential模块(其中包含我们的Linear层)，可以像前面对features一样对同样的方式对这个模块增加检查点。
如果你发现，为了让一个模型运行，你要减少批量大小，再要求得到更大的GPU之前，可以先考虑使用检查点技术！
注：文章选自《基于PyTorch的深度学习》Ian Pointer著
|