计算 Python 代码的内存和模型显存消耗的小技巧
了解Python代码的内存消耗是每一个开发人员都必须要解决的问题,这个问题不仅在我们使用pandas读取和处理CSV文件的时候非常重要,在我们使用GPU训练的时候还需要规划GPU的显存使用。尤其是我们在白嫖使用kaggle和colab时显得更为重要。
本篇文章我们将介绍两个 Python 库 memory_profiler和Pytorch-Memory-Utils,这两个库可以帮助我们了解内存和显存的消耗。
memory_profiler
Pytorch-Memory-Utils
通过Pytorch-Memory-Utils工具,我们在使用显存的代码中间插入检测函数,这样就可以输出在当前行代码时所占用的显存。这个对于我们计算模型的GPU显存占用是非常方便的,通过计算显存占用,我们才能够最大化训练的batch size,保证训练的最优速度。
Pytorch在开始运行程序时需要额外的显存开销,这种额外的显存开销与我们实际使用的模型权重显存大小无关。 这个额外的显存Pytorch的开发者也对此进行说明了,这部分释放后的显存可以用,只不过不在Nvidia-smi中显示,所以我们无需关注。
完整文章