Bootstrap

RuntimeError: CUDA out of memory. Tried to allocate XX.XX MiB. pytorch训练超出撑爆显存的问题

1、batch_size设置过大

这种比较好理解,就是单卡batch_size设置大了,数据量就大了,显存可能就放不下了。不过一般batch_size也不宜设置过小,不然如果batch里含有噪声数据其占比就会较大,对模型训练影响就比较大,有时就会把模型训飞了(亲身经历)。

如果batch_size已经调的较小了还是爆了显存,可能就是别的问题了,接着往下看。

2、中间变量多。删除清显存

一般pytorch模型动态图是写在模型类forward函数里,可以理解为显存里主要放了forward函数里的各个变量和其定义的动态计算图。所以定义动态图的时候少写中间变量,能用一行代码解决就用一行。

还有用torch.cuda.empty_cache()在一个batch训练完及时释放内存,比如

del outputs, loss                    # 先删除
torch.cuda.empty_cache()             # 再释放

这条我试了下,作用有但没那么大,治标不治本。接着往下。

3、检查梯度

还有一部分占用显存

;