如何解决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()**

总结

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