7.12 模型显存/mix-precision

最新推荐文章于 2024-09-05 17:46:58 发布

Codiplay

最新推荐文章于 2024-09-05 17:46:58 发布

阅读量338

点赞数

文章标签：深度学习人工智能机器学习

本文链接：https://blog.csdn.net/isayiwant/article/details/131478329

版权

文章讨论了模型的显存占用，包括参数量和运算量的计算，以及如何通过减少batch_size和优化模型结构来减小显存需求。另外，重点介绍了PyTorch的自动混合精度(AMP)训练，这是一种在支持TensorCore的CUDA设备上加速运算并降低显存使用的技术，通过使用半精度浮点运算和梯度缩放来保持训练精度。

摘要由CSDN通过智能技术生成

一、完全参考：模型的显存和参数量计算

显存占用=模型显存(参数)+batch_size×每个样本显存(输出和梯度动量)

首先是“运算量”和“参数量”两个概念：
参数量：这个比较好理解，例如卷积层中的卷积核c_i*k*k*n_o，其参数量就是相乘的结果。而且，无论输入图像的尺寸怎么变，只要模型结构确定，参数量就固定。还需要注意，参数都是FP32（4字节）存放，所以模型大小是参数量*4。
运算量：使用FLOPs衡量，代表浮点运算次数，这个可以衡量算法/模型的复杂度。

还要区分两个容易混淆的量：FLOPS和FLOPs：
FLOPS：Floating point Opreations Per Second，每秒浮点运算次数，理解为计算速度，是一个衡量硬件的标准。GPU算力描述的就是这个，这些数值的单位为MM=10^12次。
FLOPs：Floating point Opreations，s是复数，为浮点数操作数。理解为算法的计算量。paper中通常使用的是GFLOPs，即10亿次浮点运算。
哪些需要用到显存？

1、模型参数

参数的显存占用：只有有参数的层，才会有显存占用。这部分的显存占用和输入无关，模型加载完成之后就会占用。

有参数的层主要包括：卷积、全连接、BatchNorm、Embedding等等

（卷积神经网络的参数基本位于卷积层和全连接层）

无参数的层主要包括：多数的激活层(Sigmod/ReLU)、池化层、Dropout等等

参数占用显存=参数数目×n

n=4：float32 n=2：float16 n=8：double64

优化器如果是SGD：除了保存W之外还要保存对应的梯度，因此显存占用等于参数从占用的显存×2。如果是带Momentum-SGD，这时候还需要保存动量，因此显存×3

如果是Adam优化器，动量占用的显存更多，显存×4

故模型中与输入输出无关的显存占用包括：参数W、梯度dW(一般与参数一样)、优化器的动量

2、输入输出参数

主要看输出feature map的形状。因为输入和输出是相对的，只看一次就行

3、减少显存占用

显存占用与batch size成正比----> 降低batch_size（数据）
减少全连接层(一般只留最后一层分类用的全连接层) （模型）

二、Mix-Precision混合精度使用介绍：PyTorch的自动混合精度（AMP） - 知乎

默认的Tensor是32-bit floating point，这就是32位浮点型精度的Tensor，还有一些比如：

torch.FloatTensor (32-bit floating point)
torch.DoubleTensor (64-bit floating point)
torch.HalfTensor (16-bit floating point 1)

自动混合精度的关键词有两个：自动、混合精度：

- 混合精度预示着有不止一种精度的Tensor，那在PyTorch的AMP模块里是几种呢？2种：torch.FloatTensor和torch.HalfTensor；

- 自动预示着Tensor的dtype类型会自动变化，也就是框架按需自动调整tensor的dtype（其实不是完全自动，有些地方还是需要手工干预）；

torch.cuda.amp 的名字意味着这个功能只能在cuda上使用，事实上，这个功能正是NVIDIA的开发人员贡献到PyTorch项目中的。而只有支持Tensor core的CUDA硬件才能享受到AMP的好处（比如2080ti显卡）。Tensor Core是一种矩阵乘累加的计算单元，每个Tensor Core每个时钟执行64个浮点混合精度操作（FP16矩阵相乘和FP32累加），英伟达宣称使用Tensor Core进行矩阵运算可以轻易的提速，同时降低一半的显存访问和存储。

因此，在PyTorch中，当我们提到自动混合精度训练，我们说的就是在NVIDIA的支持Tensor core的CUDA设备上使用torch.cuda.amp.autocast （以及torch.cuda.amp.GradScaler）来进行训练

当进入autocast的上下文后，上面列出来的那些CUDA ops 会把tensor的dtype转换为半精度浮点型，从而在不损失训练精度的情况下加快运算。刚进入autocast的上下文时，tensor可以是任何类型，你不需要在model或者input上手工调用.half() ，框架会自动做，这也是自动混合精度中“自动”一词的由来。

autocast + GradScaler：

使用步骤：

1.1 首先实例化 torch.cuda.amp.autocast(enable=True) 作为上下文管理器或者装饰器，从而使脚本使用混合精度运行。注意：autocast 一般情况下只封装前向传播过程（包括loss的计算），并不包括反向传播（反向传播的数据类型与相应前向传播中的数据类型相同）

1. 2 使用Gradient scaling 防止在反向传播过程由于中梯度太小（float16无法表示小幅值的变化）从而下溢为0的情况。torch.cuda.amp.GradScaler() 可以自动进行gradient scaling。注意：由于GradScaler()对gradient进行了scale，因此每个参数的gradient应该在optimizer更新参数前unscaled，从而使学习率不受影响。

具体代码待补：

速度变慢应该有两个原因，1是单精度和半精度之间的转换开销，不过这部分开销比较小，相比之下半精度减少的后续计算量可以cover住，另一部分额外的开销应该是梯度回传时的数值放大和缩小，也就是评论中说加了scaler会变慢，这部分开销应该是蛮大的，本身需要回传的参数梯度就很多，再加上乘法和除法操作，但是如果不加scaler，梯度回传的时候就容易出现underflow（16bit能表示的精度有限，梯度值太小丢失信息会很大），所以不加scaler最后的结果可能会变差。整体来讲这是一个balance问题，属于时间换空间。