Nvidia 在Volta 架构中引入 Tensor Core 单元，来支持 FP32 和 FP16 混合精度计算。tensor core是一种矩阵乘累加的计算单元，每个tensor core 时执行64个浮点混合精度操作（FP16矩阵相乘和FP32累加）。
同年提出了一个pytorch 扩展apex，来支持模型参数自动混合精度训练。
Pytorch1.6版本以后，开始原生支持amp，即torch.cuda.amp,是 nvidia开发人员贡献到pytorch里的，只有支持 tensor core 的 CUDA 硬件才能享受到 amp 带来的优势。
1 FP16半精度
FP16 和 FP32，是计算机使用的二进制浮点数据类型。
FP16 即半精度，使用2个字节。FP32 即Float。
其中，sign为表示正负，exponent位表示指数 2 ( n − 15 + 1 ) 2^{(n-15+1)} 2(n−15+1)，具体的细节这里不说明。需要看时再百度。
float类型在内存中的表示
单独使用FP16：
- 优势：
- 减小显存的占用，从而能支持更多的batchsize、更大模型和更大的输入尺寸 进行训练，有时反而带来精度上的提升
加快训练和推理的计算，能加快一倍的速度- 缺点：
- 溢出问题：
由于FP16的动态范围比FP32 的数值范围小很多，所以在计算过程中很容易出现上溢出和下溢出，然后就出现了"NAN"的问题。在深度学习中，由于激活函数的梯度往往比权重梯度小，更易出现下溢出问题。
当第L层的梯度下溢出时，第L-1层已经以前的所有层的权重都无法更新- 舍入误差
指 当梯度过小时，小于当前区间内的最小间隔时，该次梯度更新可能失败。比如
FP16的 2 − 3 + 2 − 14 = 2 − 3 2^{-3}+2^{-14}=2^{-3} 2−3+2−14=2−3，此时就发生了舍入误差：在 [ 2 − 3 , 2 − 2 ] [2^{-3}, 2^{-2}] [2−3,2−2] 间，比 2 − 3 2^{-3} 2−3大的下一个数为 ( 2 − 3 + 2 − 13 2^{-3}+2^{-13} 2−3+2−13)
import numpy as np a = np.array(2**(-3),dtype=np.float16) b = np.array(2**(-14),dtype=np.float16) c = a+b print(a) # 0.125 print('%f'%b) # 0.000061 print(c) # 0.125
pytorch中的数据类型：
在pytorch中，一共有10中类型的tensor：
torch.FloatTensor– 32bit floating point （pytorch默认创建的tensor的类型）
torch.DoubleTensor– 64bit floating point
torch.HalfTensor– 16bit floating piont1
torch.BFloat16Tensor– 16bit floating piont2
torch.ByteTensor– 8bit integer(unsigned)
torch.CharTensor– 8bit integer(signed)
torch.ShortTensor– 16bit integer(signed)
torch.IntTensor– 32bit integer(signed)
torch.LongTensor– 64bit integer(signed)
torch.BoolTensor– Booleanimport torch tensor = torch.zeros(20,20) print(tensor.type())
2 混合精度训练机制
自动混合精度（Automatic Mixed Precision, AMP）训练，是在训练一个数值精度为32的模型时，一部分算子的操作 数值精度为FP16，其余算子的操作精度为FP32。具体的哪些算子使用的精度，amp自动设置好了，不需要用户额外设置。
这样在不改变模型、不降低模型训练精度的前提下，可以缩短训练时间，降低存储需求，从而能支持更多的batchsize、更大模型和更大的输入尺寸 进行训练。
torch.cuda.amp给用户提供了很方便的混合精度训练机制，通过使用amp.autocast和amp.GradScaler来实现：
- 用户不需要手动对模型参数的dtype，amp会自动为算子选择合适的数值精度
- 在反向传播时，FP16的梯度数值溢出的问题，amp提供了梯度scaling操作，而且在优化器更新参数前，会自动对梯度 unscaling。所以对模型优化的超参数不会有任何的影响。
具体的实现流程如下：正常的神经网络训练：前向计算loss、反向梯度计算、梯度更新。
混合精度训练：拷贝权重副本并转为FP16模型、前向计算loss、loss放大、反向梯度计算、梯度缩小、FP16 的梯度更新到FP32模型。具体的amp的训练流程：
- 维护一个 FP32 数值精度模型的副本
- 在每个迭代
- 拷贝并且转换成 FP16 模型。
- 前向传播（FP16的模型参数）
FP16的算子，直接计算操作；对 FP32 的算子，输入输出是FP16，计算的精度为FP32。反向时同理- loss 放大 s 倍
- 反向传播，也就是反向梯度计算（FP16的模型参数和参数梯度）
- 梯度乘以 1/s
- 利用 FP16 的梯度更新 FP32 的模型参数
其中放大系数 s 的选择，选择一个常量是不合适的。因为loss和梯度的数值是变化的，所以 s 需要跟着 loss 来动态变化。
健康的loss 振荡中下降，因此GradScaler设计的 s 每隔 N 个 iteration 乘一个大于1的系数，在scale loss；
- 维护一个 FP32 数值精度模型的副本
- 在每个迭代
- 拷贝并且转换成 FP16 模型。
- 前向传播（FP16的模型参数）
- loss 放大 s 倍
- 反向传播，也就是反向梯度计算（FP16的模型参数和参数梯度）
- 检查是否有
inf或者nan的参数梯度。如果有，降低 s，回到步骤1- 梯度乘以 1/s
- 利用 FP16 的梯度更新 FP32 的模型参数
用户使用混合精度训练基本操作如下：
from torch.cuda.amp import GradScaler as GradScaler # amp依赖Tensor core架构，所以model参数必须是cuda tensor类型 model = Net().cuda() optimizer = optim.SGD(model.parameters(), ...) # GradScaler对象用来自动做梯度缩放 scaler = GradScaler() for epoch in epochs: for input, target in data: optimizer.zero_grad() # 在autocast enable 区域运行forward with autocast(): # model做一个FP16的副本，forward output = model(input) loss = loss_fn(output, target) # 用scaler，scale loss(FP16)，backward得到scaled的梯度(FP16) scaler.scale(loss).backward() # scaler 更新参数，会先自动unscale梯度 # 如果有nan或inf，自动跳过 scaler.step(optimizer) # scaler factor更新 scaler.update()具体介绍如下。
3 aotucast
classs aotucast(device_type, enable=True, **kwargs)
- [device_type] (string) 表示是否使用 ‘cuda’ 或者 ‘cpu’ 设备
- [enabled] (bool，默认为True) 表示是否在区域中启用自动投射（自动转换）
- [dtype] (torch_dpython 类型) 表示使用 torch.float16/ torch.bfloat16
- [cache_enabled] (bool，默认为True) 表示是否使用 autocast 中的权重缓存
说明：
- autocast 的实例可以用作上下文管理器 或装饰器，设置区域以混合精度运行
3.1 autocast 算子
在pytorch中，在使用autocast的区域，会将部分算子自动转换成FP16 进行计算。只有CUDA算子有资格被自动转换。
- amp 自动转换成 FP16 的算子有：
- 自动转换成 FP32 的算子：
- 剩下没有列出的算子，像dot,add,cat…都是按数据中较大的数值精度，进行操作，即有 FP32 参与计算，就按 FP32，全是 FP16 参与计算，就是 FP16。
3.2 显示转换精度的情况
进入autocast-enabled 区域时，张量可以是任何类型。使用自动投射时，不应在模型或输入上调用 half() 或 bfloat16()。
但，作为上下文管理器使用时，混合精度计算enable区域得到的FP16数值精度的变量在enable区域外要显式的转换成FP32，否则使用过程中可能会导致类型不匹配的错误# 在默认数据类型中创建一些张量（此处假定为FP32） a_float32 = torch.rand((8, 8), device="cuda") b_float32 = torch.rand((8, 8), device="cuda") c_float32 = torch.rand((8, 8), device="cuda") d_float32 = torch.rand((8, 8), device="cuda") with autocast(): # torch.mm 是在 autocast算子的列表中，会转换为 FP16. # 输入为FP32, 但会以FP16精度运行计算，并输出FP16数据 # 这个过程不需要手动设置 e_float16 = torch.mm(a_float32, b_float32) # 也可以是混合输入类型 f_float16 = torch.mm(d_float32, e_float16) # 但 在退出 autocast 后，使用autocast区域生成的FP16变量时，就需要显示的转换成FP32。 g_float32 = torch.mm(d_float32, f_float16.float())
autocast也可以嵌套使用：# 在默认数据类型中创建一些张量（此处假定为FP32） a_float32 = torch.rand((8, 8), device="cuda") b_float32 = torch.rand((8, 8), device="cuda") c_float32 = torch.rand((8, 8), device="cuda") d_float32 = torch.rand((8, 8), device="cuda") with autocast(): e_float16 = torch.mm(a_float32, b_float32) with autocast(enabled=False): f_float32 = torch.mm(c_float32, e_float16.float()) g_float16 = torch.mm(d_float32, f_float32)
3.3 autocast 作为装饰器
这种情况一般使用分布式训练中。autocast 设计为 “thread local” 的，所以只在 main thread 上设 autocast 区域是不 work 的：
非分布式训练一般调用形式为：
model = MyModel() with autocast(): output = model(input)
分布式训练会用nn.DataParalle()或nn.DistributedDataParallel，在创建model之后添加相应的代码，如下，但这样是不生效的，这里的autocast只在main thread中工作：model = MyModel() DP_model = nn.DataParalle(model) ## 添加 with autocast(): output = DP_model(input)
为了在其他thread上同时也生效，需要在定义网络结构中的 forward 也设置 autocast。有两种方式，添加装饰器、添加上下文管理器。## 方式1：装饰器 class myModel(nn.Module): @autocast() def forward(self, input): pass ## 方式2：上下文管理器 class myModule(nn.Module): def forward(self, input): with autocast(): pass ## 主函数中调用 model = MyModel() DP_model = nn.DataParalle(model) ## 添加 with autocast(): output = DP_model(input)
4 GradScaler 类
当使用了混合精度训练，存在无法收敛的情况，原因是激活梯度的值太小了，造成了溢出。可以通过使用
torch.cuda.amp.GradScaler，放大loss的值 来防止梯度的underflow。
torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)
- 【init_scale】 scale factor的初始值
- 【growth_factor】 每个scale factor 的增长系数
- 【backoff_factor】scale factor 下降系数
- 【growth_interval】每隔多个interval增长scale factor
- 【enabled】是否做scale
4.1 GradScaler 的方法
scale(output)方法
对 outputs 成 scale factor，并返回。如果enabled=False，就直接返回step(optimizer, *args, **kwargs)方法
完成了两个功能：对梯度 unscale；检查梯度溢出，如果没有 nan/inf，就执行optimizer 的step，如果有就跳过update(new_scale=None)方法
update方法在每个 iteration 结束前都需要调用，如果参数更新跳过，会给scale factor 乘以 backoff_factor，或者到了该增长的iteration，就给scale factor 乘 growth_factor。也可以使用 new_scale 直接更新 scale factor.
例子：
model=Net().cuda() optimizer=optim.SGD(model.parameters(),...) scaler = GradScaler() #训练前实例化一个GradScaler对象 for epoch in epochs: for input,target in data: optimizer.zero_grad() with autocast(): ＃前后开启autocast output=model(input) loss = loss_fn(output,targt) scaler.scale(loss).backward() #为了梯度放大 #scaler.step() 首先把梯度值unscale回来，如果梯度值不是inf或NaN,则调用optimizer.step()来更新权重，否则，忽略step调用，从而保证权重不更新。 scaler.step(optimizer) scaler.update() #准备着，看是否要增大scaler
4.2 GradScaler在梯度处理更多方面的应用
Gradient clipping
scaler = GradScaler() for epoch in epochs: for input, target in data: optimizer.zero_grad() with autocast(): output = model(input) loss = loss_fn(output, target) scaler.scale(loss).backward() # 先进行unscale 梯度，此时的clip threshold才能正确对梯度使用 scaler.unscale_(optimizer) # clip梯度 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) # unscale_() 已经被显式调用了，scaler执行step时就不再unscalse更新参数，有nan/inf也会跳过 scaler.step(optimizer) scaler.update()
Gradient accumulationscaler = GradScaler() for epoch in epochs: for i, (input, target) in enumerate(data): with autocast(): output = model(input) loss = loss_fn(output, target) # loss 根据 累加的次数归一一下 loss = loss / iters_to_accumulate # scale 归一的loss 并backward scaler.scale(loss).backward() if (i + 1) % iters_to_accumulate == 0: # may unscale_ here if desired # (e.g., to allow clipping unscaled gradients) # step() and update() proceed as usual. scaler.step(optimizer) scaler.update() optimizer.zero_grad()
Gradient penaltyfor epoch in epochs: for input, target in data: optimizer.zero_grad() with autocast(): output = model(input) loss = loss_fn(output, target) # 防止溢出，在不是autocast 区域，先用scaled loss 得到 scaled 梯度 scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss), inputs=model.parameters(), create_graph=True) # 梯度unscale inv_scale = 1./scaler.get_scale() grad_params = [p * inv_scale for p in scaled_grad_params] # 在autocast 区域，loss 加上梯度惩罚项 with autocast(): grad_norm = 0 for grad in grad_params: grad_norm += grad.pow(2).sum() grad_norm = grad_norm.sqrt() loss = loss + grad_norm scaler.scale(loss).backward() # may unscale_ here if desired # (e.g., to allow clipping unscaled gradients) # step() and update() proceed as usual. scaler.step(optimizer) scaler.update()
4.5 Multiple models
只需要使用一个scaler对多个模型操作，但scale(loss) 和 step(optimizer) 要分别执行
scaler = torc h.cuda.amp.GradScaler() for epoch in epochs: for input, target in data: optimizer0.zero_grad() optimizer1.zero_grad() with autocast(): output0 = model0(input) output1 = model1(input) loss0 = loss_fn(2 * output0 + 3 * output1, target) loss1 = loss_fn(3 * output0 - 5 * output1, target) # 这里的retain_graph与amp无关，这里出现是因为在这个示例中，两个backward() 调用都共享图的一些部分。 scaler.scale(loss0).backward(retain_graph=True) scaler.scale(loss1).backward() # 如果要检查或修改其拥有的参数的梯度，可以选择相应的优化器 进行显式取消缩放。 scaler.unscale_(optimizer0) scaler.step(optimizer0) scaler.step(optimizer1) scaler.update()
5 混合精度使用注意事项
- 尽量在 具有 Tensor Core 架构的 GPU 使用amp。
在没有Tensor Core 架构的GPU 上使用amp，显存会明显减小，但速度会下降较多。具体的，在 Turing架构的 GTX 1660 上使用amp，运算时间增加了 一倍，显存不到原来的一半- 常数范围：为了保证计算不溢出，首先保证人工设定的常数不溢出。如 epsilon、INF 等
- Dimension 最好是8的倍数：维度是8的倍数，性能最好
5472
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
