在 PyTorch 中使用梯度检查点在GPU 上训练更大的模型

作为机器学习从业者，我们经常会遇到这样的情况，想要训练一个比较大的模型，而 GPU 却因为内存不足而无法训练它。当我们在出于安全原因不允许在云计算的环境中工作时，这个问题经常会出现。在这样的环境中，我们无法足够快地扩展或切换到功能强大的硬件并训练模型。并且由于梯度下降算法的性质，通常较大的批次在大多数模型中会产生更好的结果，但在大多数情况下，由于内存限制，我们必须使用适应GPU显存的批次大小。

本文将介绍解梯度检查点（Gradient Checkpointing），这是一种可以让你以增加训练时间为代价在 GPU 中训练大模型的技术。 我们将在 PyTorch 中实现它并训练分类器模型。

完整文章：

https://avoid.overfit.cn/post/a13e29c312c741ac94d4a5079fb9f8af