python

关注公众号 jb51net

关闭
首页 > 脚本专栏 > python > pytorch训练过程中CPU内存溢出问题

如何解决pytorch训练过程中CPU内存溢出问题

作者:yyyyyyhx

这篇文章主要介绍了如何解决pytorch训练过程中CPU内存溢出问题,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教

pytorch训练过程中CPU内存溢出问题

期待了一晚上的结果,早上起来竟然发现CPU内存溢出导致程序奔溃了,今天郁闷了一天。。。

经查询,内存溢出的常见原因有:

排查过程

找了好多原因,感觉上面的比较贴近,但是改了一堆,内存还是蹭蹭往上涨。

后来调用了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()**

总结

以上为个人经验,希望能给大家一个参考,也希望大家多多支持脚本之家。

您可能感兴趣的文章:
阅读全文