前言
Python 的 memory_profiler这个库可以帮助我们了解内存的消耗,
Pytorch-Memory-Utils和pynvml这个库可以帮助我们了解显存消耗
计算显存和内存
具体用法参考:
参考文档:https://blog.csdn.net/yrwang_xd/article/details/105949077
参考文档:https://www.overfit.cn/post/aecb96d60db74f609133181899f172cb
参考文档:https://oldpan.me/archives/how-to-calculate-gpu-memory
优化
-
去掉全连接层,改为全局平均池化
-
降低batch size
-
池化,减小特征图size
-
pytorch中的inplace操作,torch.no_grad()
参考文档:https://oldpan.me/archives/how-to-use-memory-pytorch
参考文档:https://zhuanlan.zhihu.com/p/31558973