CNN卷积神经网络模型的GPU显存占用分析

一、参考资料

浅谈深度学习:如何计算模型以及中间变量的显存占用大小
 如何在Pytorch中精细化利用显存

二、相关介绍

0. 预备知识

为了方便计算，本文按照以下标准进行单位换算：

1 G = 1000 MB
1 M = 1000 KB
1 K = 1000 Byte
1 B = 8 bit

1. 模型参数量的计算方法

参考博客：CNN卷积神经网络模型的参数量、计算量计算方法（概念版）

2. 张量的数据类型

Data type	dtype	CPU tensor	GPU tensor
32-bit floating point	`torch.float32` or `torch.float`	`torch.FloatTensor`	`torch.cuda.FloatTensor`
64-bit floating point	`torch.float64` or `torch.double`	`torch.DoubleTensor`	`torch.cuda.DoubleTensor`
16-bit floating point	`torch.float16` or `torch.half`	`torch.HalfTensor`	`torch.cuda.HalfTensor`
8-bit integer (unsigned)	`torch.uint8`	`torch.ByteTensor`	`torch.cuda.ByteTensor`
8-bit integer (signed)	`torch.int8`	`torch.CharTensor`	`torch.cuda.CharTensor`
16-bit integer (signed)	`torch.int16` or `torch.short`	`torch.ShortTensor`	`torch.cuda.ShartTensor`
32-bit integer (signed)	`torch.int32` or `torch.int`	`torch.IntTensor`	`torch.cuda.IntTensor`
64-bit integer (signed)	`torch.int64` or `torch.long`	`torch.LongTensor`	`torch.cuda.LongTensor`

通常，模型训练使用以下两种数据类型：

float32 单精度浮点型；
int32 整型。

8bit的整型int所占的空间为 1B，32bit的浮点型float所占空间为 4B。而double双精度浮点型和长整型long在平常的模型训练中一般不会使用。

消费级显卡对单精度计算有优化，服务器级显卡对双精度计算有优化。

3. 关于`inplace=False`

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

三、显存占用相关介绍

0. 引言

torch.FatalError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1524590031827/work/aten/src/THC/generic/THCStorage.cu:58

由于显存不足，导致程序崩溃。学会计算模型以及中间变量所占显存大小，是很有必要的。

1. 图片的显存占用

假设一张RGB三通道真彩色图片，长宽分别为 500x500，数据类型为单精度浮点型，那么这张图片所占显存的大小为：500x500x3x4B=3MB。而一个 (256, 3, 100, 100)-(N, C, H, W) 的FloatTensor所占的空间为：256x3x100x100x4B = 31MB

2. 模型的显存占用

通常，模型占用的显存来自两个部分：

模型自身的参数（params），即有参数的网络层。
模型在计算时产生的中间参数（memory）。

一般来说，模型自身参数并不会占用很多的显存空间，主要占用显存空间的是计算时产生的中间参数。

2.1 模型自身的参数（params）

有参数的网络层，包括：

卷积层：Conv2d(Cin, Cout, K)，参数量为 Cin × Cout × K × K
全连接层：Linear(M->N)，参数量为 M×N
BatchNorm层：BatchNorm(N)，参数量为 2N
Embedding层：Embedding(N,W)，参数量为 N × W

不带参数的网络层，包括：

激活层Relu等；
池化层；
Dropout层；

2.2 模型的中间参数（memory）

模型在计算时产生的中间参数，也就是输入图像在计算时每一层产生的输入和输出；
backward反向传播计算时产生的额外的中间参数；
优化器在优化时产生的额外的模型参数。

3. 实际显存与理论显存

为什么实际占用的显存空间比理论计算的大？
大概原因是深度学习框架一些额外的开销。不过，通过上面公式计算出来的显存理论值和实际值不会相差太多。

4. 计算显存占用大小

4.1 方法一（推荐）

使用 torchstat 工具，计算模型的显存占用大小。参考博客：CNN卷积神经网络模型的参数量计算方法（经验版）

4.2 方法二

当然，也可以自定义函数计算显存占用大小，代码如下所示：

# 模型显存占用监测函数
# model：输入的模型
# input：实际中需要输入的Tensor变量
# type_size 默认为 4 默认类型为 float32 

def modelsize(model, input, type_size=4):
    para = sum([np.prod(list(p.size())) for p in model.parameters()])
    print('Model {} : params: {:4f}M'.format(model._get_name(), para * type_size / 1000 / 1000))

    input_ = input.clone()
    input_.requires_grad_(requires_grad=False)

    mods = list(model.modules())
    out_sizes = []

    for i in range(1, len(mods)):
        m = mods[i]
        if isinstance(m, nn.ReLU):
            if m.inplace:
                continue
        out = m(input_)
        out_sizes.append(np.array(out.size()))
        input_ = out

    total_nums = 0
    for i in range(len(out_sizes)):
        s = out_sizes[i]
        nums = np.prod(np.array(s))
        total_nums += nums


    print('Model {} : intermedite variables: {:3f} M (without backward)'
          .format(model._get_name(), total_nums * type_size / 1000 / 1000))
    print('Model {} : intermedite variables: {:3f} M (with backward)'
          .format(model._get_name(), total_nums * type_size*2 / 1000 / 1000))

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

5. 显存优化方法

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

优化除了算法层的优化，最基本的显存优化无非也就以下几点：

减少输入图像的尺寸；
减少batch，减少每次的输入图像数量；
多使用下采样，池化层；
一些神经网络层可以进行小优化，利用relu层中设置inplace；
购买显存更大的显卡；
从深度学习框架上面进行优化。

5.1 牺牲计算速度减少显存使用量

在PyTorch中，如果一个模型占用的显存太大了，可以将一个计算过程分成两半，先计算前一半，保存后一半需要的中间结果，然后再计算后一半。

# 首先设置输入的input=>requires_grad=True
# 如果不设置可能会导致得到的gradient为0

# 输入
input = torch.rand(1, 10, requires_grad=True)

# 假设我们有一个非常深的网络
layers = [nn.Linear(10, 10) for _ in range(1000)]


# 定义要计算的层函数，可以看到我们定义了两个
# 一个计算前500个层，另一个计算后500个层

def run_first_half(*args):
    x = args[0]
    for layer in layers[:500]:
        x = layer(x)
    return x

def run_second_half(*args):
    x = args[0]
    for layer in layers[500:-1]:
        x = layer(x)
    return x

# 我们引入新加的checkpoint
from torch.utils.checkpoint import checkpoint

x = checkpoint(run_first_half, input)
x = checkpoint(run_second_half, x)

# 最后一层单独调出来执行
x = layers[-1](x)
x.sum.backward()

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

input = torch.rand(1, 10, requires_grad=True)
layers = [nn.Linear(10, 10) for _ in range(1000)]
model = nn.Sequential(*layers)

from torch.utils.checkpoint import checkpoint_sequential

# 分成两个部分
num_segments = 2
x = checkpoint_sequential(model, num_segments, input)
x.sum().backward()

6. 跟踪显存使用情况

再次浅谈Pytorch中的显存利用问题(附完善显存跟踪代码)

我们借用 Pytorch-Memory-Utils 这个工具来检测我们在训练过程中关于显存的变化情况，分析出我们如何正确释放多余的显存。

通过Pytorch-Memory-Utils工具，我们在使用显存的代码中间插入检测函数，就可以输出类似于下面的信息，At __main__ <module>: line 13 Total Used Memory:696.5 Mb表示在当前行代码时所占用的显存，即在我们的代码中执行到13行的时候所占显存为695.5Mb。At __main__ <module>: line 15 Total Used Memory:1142.0 Mb表示程序执行到15行时所占的显存为1142.0Mb。两条数据之间表示所占显存的tensor变量。

# 12-Sep-18-21:48:45-gpu_mem_track.txt

GPU Memory Track | 12-Sep-18-21:48:45 | Total Used Memory:696.5  Mb

At __main__ <module>: line 13                        Total Used Memory:696.5  Mb

+ | 7 * Size:(512, 512, 3, 3)     | Memory: 66.060 M | <class 'torch.nn.parameter.Parameter'>
+ | 1 * Size:(512, 256, 3, 3)     | Memory: 4.7185 M | <class 'torch.nn.parameter.Parameter'>
+ | 1 * Size:(64, 64, 3, 3)       | Memory: 0.1474 M | <class 'torch.nn.parameter.Parameter'>
+ | 1 * Size:(128, 64, 3, 3)      | Memory: 0.2949 M | <class 'torch.nn.parameter.Parameter'>
+ | 1 * Size:(128, 128, 3, 3)     | Memory: 0.5898 M | <class 'torch.nn.parameter.Parameter'>
+ | 8 * Size:(512,)               | Memory: 0.0163 M | <class 'torch.nn.parameter.Parameter'>
+ | 3 * Size:(256, 256, 3, 3)     | Memory: 7.0778 M | <class 'torch.nn.parameter.Parameter'>
+ | 1 * Size:(256, 128, 3, 3)     | Memory: 1.1796 M | <class 'torch.nn.parameter.Parameter'>
+ | 2 * Size:(64,)                | Memory: 0.0005 M | <class 'torch.nn.parameter.Parameter'>
+ | 4 * Size:(256,)               | Memory: 0.0040 M | <class 'torch.nn.parameter.Parameter'>
+ | 2 * Size:(128,)               | Memory: 0.0010 M | <class 'torch.nn.parameter.Parameter'>
+ | 1 * Size:(64, 3, 3, 3)        | Memory: 0.0069 M | <class 'torch.nn.parameter.Parameter'>

At __main__ <module>: line 15                        Total Used Memory:1142.0 Mb

+ | 1 * Size:(60, 3, 512, 512)    | Memory: 188.74 M | <class 'torch.Tensor'>
+ | 1 * Size:(30, 3, 512, 512)    | Memory: 94.371 M | <class 'torch.Tensor'>
+ | 1 * Size:(40, 3, 512, 512)    | Memory: 125.82 M | <class 'torch.Tensor'>

At __main__ <module>: line 21                        Total Used Memory:1550.9 Mb

+ | 1 * Size:(120, 3, 512, 512)   | Memory: 377.48 M | <class 'torch.Tensor'>
+ | 1 * Size:(80, 3, 512, 512)    | Memory: 251.65 M | <class 'torch.Tensor'>

At __main__ <module>: line 26                        Total Used Memory:2180.1 Mb

- | 1 * Size:(120, 3, 512, 512)   | Memory: 377.48 M | <class 'torch.Tensor'> 
- | 1 * Size:(40, 3, 512, 512)    | Memory: 125.82 M | <class 'torch.Tensor'> 

At __main__ <module>: line 32                        Total Used Memory:1676.8 Mb

当然这个检测工具不仅适用于Pytorch，其他的深度学习框架也同样适用，不过需要注意下静态图和动态图在实际运行过程中的区别。