PyTorch 节省显存的策略总结

文章预览

点击上方 “ 小白学视觉 ”，选择加" 星标 "或“ 置顶 ” 重磅干货，第一时间送达 作者丨 OpenMMLab 来源丨 https://zhuanlan.zhihu.com/p/430123077 编辑丨极市平台 极市导读   随着深度学习快速发展，同时伴随着模型参数的爆炸式增长，对显卡的显存容量提出了越来越高的要求，如何在单卡小容量显卡上面训练模型是一直以来大家关心的问题。本文结合 MMCV 开源库对一些常用的节省显存策略进行了简要分析。   0 前言 本文涉及到的 PyTorch 节省显存的策略包括： 混合精度训练 大 batch 训练或者称为梯度累加 gradient checkpointing 梯度检查点 1 混合精度训练 混合精度训练全称为 Automatic Mixed Precision，简称为 AMP，也就是我们常说的 FP16。在前系列解读中已经详细分析了 AMP 原理、源码实现以及 MMCV 中如何一行代码使用 AMP，具体链接见： OpenMMLab：PyTorch 源码解读之 torch.c ………………………………