终于解决pytorch训练过程中CPU内存溢出的问题
期待了一晚上的结果,早上起来竟然发现CPU内存溢出导致程序奔溃了,今天郁闷了一天。。。
经查询,内存溢出的常见原因有:
loss求和未加item()
num_workers过大
大量使用list转tensor
排查过程
找了好多原因,感觉上面的比较贴近,但是改了一堆,内存还是蹭蹭往上涨。后来调用了memory_profiler这个包,发现了
我的程序中上面这部分每轮都加了70多m,可能是这个循环导致的(另外排查了好多程序都不说了。。。)
然后我把这个循环改成了下面这样
观察每组数据的memory占用,发现了很神奇的现象:
第一组数据没怎么占内存,后面重复调用同一个数据的时候内存就开始暴涨,到现在也没整明白为什么会这样。。欢迎大神们指教
解决方法
后来调用了
from einops import rearrange
把数组维度修改了一下,再一起送进网络
可算解决了