在一次测试一个超分辨模型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()提供一个上下文, 阻止对梯度的计算. 试了一下, 不能解决问题.

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

《南溪的目标检测学习笔记》——训练PyTorch模型遇到显存不足的情况怎么办(“OOM: CUDA out of memory“)
1 前言在目标检测中,可能会遇到显存不足的情况,我们在这里记录一下解决方案: 2 如何判断真正是出现显存溢出(不是"软件误报") 当前需要分配的显存在600MiB以下, 例如: ...
pytorch训练时显存溢出
网络在前期可以正常训练,但训练几轮后就发生显存爆炸的问题,调整输入大小或者每次循环都清除显存也无法解决问题,后来经过查询,是在对loss求和时,直接使用 tl += loss 可以看到,loss是张 ...
pytorch 优化GPU显存占用，避免out of memory
pytorch 优化GPU显存占用,避免out of memory 分享一个最实用的招: 用完把tensor删掉,pytorch不会自动清理显存! 代码举例,最后多删除一个,gpu显存占用就会下降,训 ...
解决Tensorflow显存溢出的问题
import tensorflow as tfimport tensorflow.keras.backend as Kos.environ["TF_FORCE_GPU_ALLOW_GROWT ...
模型的显存和参数量计算
写在前面:以此记录关于模型显存和参数量的一些理解和计算. 首先是"运算量"和"参数量"两个概念: 参数量:这个比较好理解,例如卷积层中的卷积核c_i*k*k*n ...
使用FVD时遇到的问题记录：使用tensorflow时batch_size不变的情况下使用数据量更大的dataset发生了显存溢出(OOM)问题的解决
情况:计算 FVD 时,使用 dataset-1 时运行成功,使用数据量(data size)更大的 dataset-2 时运行失败.报错OOM:显存溢出. 查找问题:比较两次计算时的 input v ...
解决测试阶段显存增大问题
转自简书:解决 pytorch 在训练时由于设置了验证集导致 out of memory (同样可用于测试时减少显存占用) 问题描述在跑pytorch的时候,在训练阶段监控到显存占用2.7G左右,但 ...
亚马逊和脸书发布 TorchServe 解决Pytorch 模型的部署问题
4月21日官方发布,PyTorch 用户可以使用TorchServe(PyTorch生态新的模型服务框架)来大规模部署经过训练的模型,而无需编写自定义代码. PyTorch是最初由Facebook创建 ...
pytorch测试模型时显存不够的问题
在训练过程中使用了分块训练,但是在测试的时候使用了整张图片输入到网络中,然后就出现了显存不够的问题经验告诉我不应该不够,但是使用了分块训练,我也不知道是不是真的不够一筹莫展之时,发现没有使用tor ...

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

分析

解决方法

注意

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

最新文章

热门文章