显存占用分析

显存分析

数据类型

深度学习框架中, 可以使用各种各样的数据类型, 数据类型的命名规范为Type + Num, 如int64, float32等. 其中的Num表示该类型所占据的比特(bit)数目.

常用的数据类型占用的大小为:

类型	bit大小	字节大小
int8	8	1Byte(1个字节)
int16	16	2Byte
int32	32	4Byte
int64	64	8Byte
float32	32	4Byte
float64	64	8Byte

float32作为深度学习中最常用的数值类型, 称为单精度浮点数, 占用4个字节. 对应的, float64称为双精度浮点数.

如果一个参数矩阵的大小为$1000 \times 1000$, 使用float32精度, 则占用的显存大小约为:

$1000 \times 1000 \times 4 \text{Byte} = 4 \text{MB}$

显存占用来源

直观看, 占用显存的内容包括:

• 模型参数

• 模型每一层的输出

除此之外, 优化器要计算并保存模型参数的梯度, 一般来说占用的大小与模型参数大小相同, 因此占用显存的来源又多了:

• 优化器中的模型参数梯度

现在使用的优化器, 如Adam等, 既要利用动量的惯性, 加快收敛, 跳出震荡, 又要使用二阶动量对每个参数的学习率实现自调节. 而这些也是要一直保存, 不断更新的, 各自占用的大小与梯度的大小相同. 因此又有占用来源:

• 优化器中的模型参数的一阶动量和二阶动量

导致显存占用较多的原因

按模型参数和模型每一层的输出分两部分讨论.

模型参数

所有的优化器都要计算并保存梯度, 以模型参数占用的显存大小为单位, 因此纯净的SGD占用2个单位的显存, 增加动量的momentum-SGD占用3个单位的显存, 像Adam这样的优化器要占用4个单位的显存.

因此在训练时, 如果显存资源紧缺, 优化器也要仔细选择. 考虑Adam相对于SGD使用方便, 不需要细致的炼丹调参, 偏向使用Adam这种更复杂的优化器. 但如果受限与显存大小, 也可以使用例如AdaFactor这种稍微节省显存的优化器, 基本对最优效果没有影响.

模型每一层的输出

之所以要保存每一层的输出, 是因为在反向传播计算参数的梯度时, 需要用到本层的输出结果. 如果模型层数较多, 比较复杂, 例如BERT, 这部分就会占用很多的显存.

这部分的一个优化方法是使用重计算, 简单来说就是用时间换空间, 在前向时只保存部分中间节点结果, 在反向时重新计算没保存的部分, 用来给计算梯度使用.

这个技巧在标准的BERT上很好用, 能够显著的增大Batch size.

常见层资源消耗分析

分析常见层的显存占用和需要的计算量, 其中显存的占用量:

显存占用 = 模型显存占用 + batch_size × 每个样本的显存占用

Linear

假设输入的feature map的hidden size为$M$, 输出的hidden size为$N$, $B$为batch size大小.

参数占用的大小为$M \times N$.

输出的feature map大小为$B \times N$.

因此对于Dense层来说, 参数占用的显存会很大, 因此尽量减少Dense层的使用.

计算量为$B \times M \times N$.

Conv1D

这里考虑在序列中使用的一维卷积核, 图像中使用的二维卷积核类推.

设卷积核的大小为$K$, 输入feature map的通道数为$C_{in}$, 输出的feature map的通道数为$C_{out}$.

参数占用的大小为$C_{in} \times C_{out} \times K$. 参数占用的显存是很小的.

输出的feature map大小为$B \times S \times C_{out}$

计算量为$B \times S \times C_{in} \times C_{out} \times K$

Transformer

首先说明这里说的Transformer指的是的Bert模型中的一层, 完整的子层顺序如下:

Multihead Attention -> Add -> Layer Normalization -> Feed Forward -> Add -> Layer Normalization

这里说明Multihead Attention, Feed Forward层的情况. 假设序列的长度为$S$, 输入输出的feature map的大小都为$H$.

Multihead Attention

Multihead Attention中的参数就是对三路输入Q, K, V以及输出O进行线性转换及激活的4个Dense层中的参数, 因此对应的大小为$4 \times H \times H$, 其中输入输出的大小都为$H$.

计算量见self-attention的时间复杂度.

Feed Forward

Feed Forward可以看做两层Dense, 先将输入的hidden size扩大, 然后再缩小为原来的大小. 因此假设中间的hidden size大小为$I$, 参数大小则为$2 \times H \times I$

总的来看, Transformer结构中大量使用了Dense结构, 再加上MultiHead的结构也增大了计算量, 整体使用的显存和需要的计算量都是很大的.

节省显存的方法

列举一些常用的方法.

减小Batch size. 减小Batch size往往都是受显存限制的无奈之举, 基本都会带来模型性能的损失. 可以考虑使用梯度累积方法, 用时间换空间.

仔细选择优化器. 如果显卡资源特别拮据, 就只能用SGD了. 另外可以尝试使用一些节省显存的Adam变种优化器, 例如AdaFactor.

重计算技巧. 用空间换时间, 在前向时只保存部分中间节点结果, 在反向时重新计算没保存的部分.

减少全连接层. Dense层的参数矩阵是非常大的. 一般只在最后的输出层使用.

精简词典. Word Embedding矩阵也是一个非常大的矩阵. BERT模型的预训练参数都会带有一个字典, 如果部分token用不到可以精简掉, 缩小Embedding矩阵的大小. 这一点在使用一些multi-language模型时经常使用.

参考资料

• 科普帖：深度学习中GPU和显存分析