【磕岩日记】模型搭建和训练过程中的显存管理

问题-“显存泄漏”

问题描述：

训练模型时，发现随着训练次数的增加，显存占用逐渐增大，直到OOM

问题原因：

1-register buffer或register module函数需要放在with torch no grad上下文中

参考：【pytorch笔记，训练时显存一直增加到out-of-memory?真相了！-AI码上来】

2-将同样的张量多次通过to('cuda')移动到GPU上

解释：这样会在GPU上创造多个副本，而且不会自动释放

3-推理时，先将结果移动到CPU上，再保存

解释：否则显存会被逐渐累积的输出数据占用

4-【特定】搭建的模型用到了一些预训练模型，可以把预训练模型放到cpu上，降低问题复杂度

新学到的技巧：

dataloader默认处于CPU上，如果有需要，可以显式地to('cuda')

设置梯度积累步数：gradient accumulation step：batch==1, gradient_accumulate_step==2, 等价于batch==2, gradient_accumulate_step==1，据说可以节省显存，但实验表明没区别

疑问：

推理时的显存占用（22280M）比训练时（21594M）还大一点点（686M）

如何减少推理时的显存占用？