如何解决pytorch训练过程中CPU内存溢出问题
作者:yyyyyyhx
这篇文章主要介绍了如何解决pytorch训练过程中CPU内存溢出问题,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教
pytorch训练过程中CPU内存溢出问题
期待了一晚上的结果,早上起来竟然发现CPU内存溢出导致程序奔溃了,今天郁闷了一天。。。
经查询,内存溢出的常见原因有:
- loss求和未加item()
- num_workers过大
- 大量使用list转tensor
排查过程
找了好多原因,感觉上面的比较贴近,但是改了一堆,内存还是蹭蹭往上涨。
后来调用了memory_profiler这个包,发现了
我的程序中上面这部分每轮都加了70多m,可能是这个循环导致的(另外排查了好多程序都不说了。。。)
然后我把这个循环改成了下面这样
观察每组数据的memory占用,发现了很神奇的现象:
第一组数据没怎么占内存,后面重复调用同一个数据的时候内存就开始暴涨,到现在也没整明白为什么会这样。。
解决方法
后来调用了
from einops import rearrange
把数组维度修改了一下,再一起送进网络
可算解决了
pytorch内存溢出,Ubuntu进程killed问题
pytorch显存越来越多的一个原因
optimizer.zero_grad() loss.backward() optimizer.step() train_loss += loss
参考了别人的代码发现那句loss一般是这样写
loss_sum += loss.data[0]
这是因为输出的loss的数据类型是Variable。而PyTorch的动态图机制就是通过Variable来构建图。
主要是使用Variable计算的时候,会记录下新产生的Variable的运算符号,在反向传播求导的时候进行使用。
如果这里直接将loss加起来,系统会认为这里也是计算图的一部分,也就是说网络会一直延伸变大那么消耗的显存也就越来越大。
用Tensor计算要写成:
train_loss += loss.item()
correct_total += torch.eq(predict, label_batch).sum().item() train_loss += loss.item()
当需要将模型中变量提取出来参与计算时,需要使用** .item()**
总结
以上为个人经验,希望能给大家一个参考,也希望大家多多支持脚本之家。