在上篇文章《浅谈深度学习:如何计算模型以及中间变量的显存占用大小》中我们对如何计算各种变量所占显存大小进行了一些探索。而这篇文章我们着重讲解如何利用Pytorch深度学习框架的一些特性，去查看我们当前使用的变量所占用的显存大小，以及一些优化工作。以下代码所使用的平台框架为Pytorch。

在Pytorch中优化显存是我们处理大量数据时必要的做法，因为我们并不可能拥有无限的显存。显存是有限的，而数据是无限的，我们只有优化显存的使用量才能够最大化地利用我们的数据，实现多种多样的算法。

上篇文章中说过，一个模型所占的显存无非是这两种：

其实权重参数一般来说并不会占用很多的显存空间，主要占用显存空间的还是计算时产生的中间变量，当我们定义了一个model之后，我们可以通过以下代码简单计算出这个模型权重参数所占用的数据量：

假设我们有这样一个model：

然后我们得到的para是112576，但是我们计算出来的仅仅是权重参数的“数量”，单位是B，我们需要转化一下：

这样就可以打印出：

但是我们之前说过一个神经网络的模型，不仅仅有权重参数还要计算中间变量的大小。怎么去计算，我们可以假设一个输入变量，然后将这个输入变量投入这个模型中，然后我们主动提取这些计算出来的中间变量：

上面得到的值是模型在运行时候产生所有的中间变量的“数量”，当然我们需要换算一下：

因为在backward的时候所有的中间变量需要保存下来再来进行计算，所以我们在计算backward的时候，计算出来的中间变量需要乘个2。

然后我们得出，上面这个模型的中间变量需要的占用的显存，很显然，中间变量占用的值比模型本身的权重值多多了。如果进行一次backward那么需要的就更多。

我们总结一下之前的代码：

当然我们计算出来的占用显存值仅仅是做参考作用，因为Pytorch在运行的时候需要额外的显存值开销，所以实际的显存会比我们计算的稍微大一些。

我们都知道激活函数Relu()有一个默认参数inplace，默认设置为False，当设置为True时，我们在通过relu()计算时的得到的新值不会占用新的空间而是直接覆盖原来的值，这也就是为什么当inplace参数设置为True时可以节省一部分内存的缘故。

在Pytorch-0.4.0出来了一个新的功能，可以将一个计算过程分成两半，也就是如果一个模型需要占用的显存太大了，我们就可以先计算一半，保存后一半需要的中间结果，然后再计算后一半。

也就是说，新的checkpoint允许我们只存储反向传播所需要的部分内容。如果当中缺少一个输出(为了节省内存而导致的)，checkpoint将会从最近的检查点重新计算中间输出，以便减少内存使用(当然计算时间增加了)：

上面的模型需要占用很多的内存，因为计算中会产生很多的中间变量。为此checkpoint就可以帮助我们来节省内存的占用了。

对于Sequential-model来说，因为Sequential()中可以包含很多的block，所以官方提供了另一个功能包：

显存的使用情况，在编写程序中我们可能无法精确计算，但是我们可以通过pynvml这个Nvidia的Python环境库和Python的垃圾回收工具，可以实时地打印我们使用的显存以及哪些Tensor使用了我们的显存。

类似于下面的报告：

以下是相关的代码，目前代码依然有些地方需要修改，等修改完善好我会将完整代码以及使用说明放到github上：https://github.com/Oldpan/Pytorch-Memory-Utils 请大家多多留意。