使用VGG网络训练发生错误RuntimeError: CUDA out of memory解决方案:

Posted 2023-02-21 算法与编程之美

问题

在使用VGG网络训练Mnisist数据集时，发生错误RuntimeError: CUDA out of memory. Tried to allocate 392.00 MiB (GPU 0; 2.00 GiB total capacity; 1.45 GiB already allocated; 0 bytes free; 1.47 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF。

网络

class VGG(nn.Module):  # 适用于(128,3,224,224)
def __init__(self):
super().__init__()
self.net = nn.Sequential(
           nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, stride=1, padding=1),
           nn.ReLU(inplace=True),
           nn.MaxPool2d(2),
           nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
           nn.MaxPool2d(2),
           nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
           nn.MaxPool2d(2),
           nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),
           nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
           nn.MaxPool2d(2),
           nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1),
           nn.MaxPool2d(2),
           nn.Flatten(),
           nn.Linear(in_features=512 * 7 * 7, out_features=4096),
           nn.Linear(in_features=4096, out_features=4096),
           nn.Linear(in_features=4096, out_features=1000),
           nn.Linear(in_features=1000, out_features=10),

       )

def forward(self, x):
       out = self.net(x)
return out

报错截图：

方法

尝试1 关闭显卡占用

根据报错（CUDA out of memory.），说明显卡内存不够。于是进入终端查一下memory现在的状态。没有在运行的进程，运行程序错误仍然存在。

尝试2 定时清理内存

在每个训练周期处插入以下代码（定时清内存）：

import torch, gc

for epoch in range(100):

...............

gc.collect()

torch.cuda.empty_cache()

尝试3 设置锁页内存pin_memory

pin_memory就是锁页内存，创建DataLoader时，设置pin_memory=True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。

主机中的内存，有两种存在方式，一是锁页，二是不锁页，锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘），而不锁页内存在主机内存不足时，数据会存放在虚拟内存中。显卡中的显存全部是锁页内存,当计算机的内存充足的时候，可以设置pin_memory=True。当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。因为pin_memory与电脑硬件性能有关，pytorch开发者不能确保每一个人都有高端设备，因此pin_memory默认为False。

尝试4 修改batch_size (问题解决)

最简单最直接的方法就是修改batch_size的大小，从而降低对显卡内存的占用。当将batch_size=4时，程序成功运行。

总结

当网路变得复杂的时候，对计算机硬件资源的要求也会随之变高。在本实验之中，就出现了因为VGG网络模型较为复杂，对计算机GPU的资源要求和消耗也随之变大，于是便出现了CUDA out of memory.GPU内存不够而报错的情况。在计算机硬件资源有限的情况下，只有选择降低batch_size的大小，从而达到计算机处理数据的能力之类。但是当batch_size过低训练出来的模型也会因为实际模型的需求出现一定的问题。简而言之batch size过小，需要花费更多时间，同时梯度震荡严重，不利于收敛；batch size过大，不同batch的梯度方向没有任何变化，容易陷入局部极小值。

这里提供一个免费使用高性能GPU的路径，在计算机资源有限前提下，如果能够使用上Google服务器可以将代码上传到Google Calab(Google colab是一个免费的 Jupyter 笔记本环境，不需要进行任何设置就可以使用，并且完全在云端运行。并且可以免费使用Google的GPU)训练，再将训练好结果下载。