计算 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中显示，所以我们无需关注。

 

完整文章

计算 Python 代码的内存和模型显存消耗的小技巧