不优雅地解决pytorch模型测试阶段显存溢出问题

在一次测试一个超分辨模型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张量 和 用于后向传播中求导的中间值.

解决方法

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

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

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

注意

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