在一次测试一个超分辨模型LESRCNN(作者提供了已训练好的模型)时,发生了CUDA out of memory的错误(虽然显卡有8G显存,但还是差了些):

RuntimeError: CUDA out of memory. Tried to allocate 1.54 GiB (GPU 0; 7.79 GiB total capacity; 4.29 GiB already allocated; 1.29 GiB free; 31.68 MiB cached)

模型的forward函数的部分代码如下图所示.

分析

模型加载后占用显存457M,但调试时发现,在forward中每执行一步都会增加对显存的占用. 占用的空间主要用处由二:计算产生的torch张量和用于后向传播中求导的中间值.

解决方法

若产生的张量后续不再被使用, 则需将其释放, 方法是使用del关键字. 如下, 发现在执行第125行时显存溢出, 于是在此前现将x等后续不再使用的张量释放.should_del是我给forward函数加的参数, 在模型测试阶段将其设为True.

使用torch.utils.checkpoint函数. 该函数“以算力换取内存”(trade compute for memory):在前向传播时不保存用于反向传播的激活值,而是在反向传播时再重新计算激活值.如下,net为模型实例, 而checkpoint函数剩余的参数为传递给模型的参数, 即模型的forward函数的参数.sr为模型的输出.

sr = torch.utils.checkpoint.checkpoint(net, lr, cfg.scale)

注意

在del后通过nvidia-smi指令查看显卡状况时会看到程序占用的显存依然不减少, 但实际上其中部分显存已经释放, 后续在该程序中新分配的张量会在这部分已释放的显存中取用.有说法说用with torch.no_grad()提供一个上下文, 阻止对梯度的计算. 试了一下, 不能解决问题.