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:continueout = m(input_)out_sizes.append(np.array(out.size()))input_ = outtotal_nums = 0for i in range(len(out_sizes)):s = out_sizes[i]nums = np.prod(np.array(s))total_nums += numsprint('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 xdef run_second_half(*args):x = args[0]for layer in layers[500:-1]:x = layer(x)return x# 我们引入新加的checkpoint
from torch.utils.checkpoint import checkpointx = 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.txtGPU Memory Track | 12-Sep-18-21:48:45 | Total Used Memory:696.5  MbAt __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，其他的深度学习框架也同样适用，不过需要注意下静态图和动态图在实际运行过程中的区别。