简介
AMP:Automatic mixed precision
,自动混合精度,可以在神经网络推理过程中,针对不同的层,采用不同的数据精度进行计算,从而实现节省显存和加快速度的目的。在Pytorch 1.5
版本及以前,通过NVIDIA提供的apex库可以实现amp
功能。但是在使用过程中会伴随着一些版本兼容和奇怪的报错问题。从1.6版本开始,Pytorch
原生支持自动混合精度训练,并已进入稳定阶段,AMP
训练能在 Tensor Core GPU
上实现更高的性能并节省多达 50% 的内存。
环境
- Python 3.8
- Pytorch 1.7.1
- CUDA 11 + cudnn 8
- NVIDIA GeFore RTX 3070
原理
关于低精度计算
当前的深度学习框架大都采用的都是FP32
来进行权重参数的存储,比如Python float
的类型为双精度浮点数 FP64
,PyTorch Tensor
的默认类型为单精度浮点数FP32
。随着模型越来越大,加速训练模型的需求就产生了。在深度学习模型中使用FP32
主要存在几个问题,第一模型尺寸大,训练的时候对显卡的显存要求高;第二模型训练速度慢;第三模型推理速度慢。其解决方案就是使用低精度计算对模型进行优化。
推理过程中的模型优化目前比较成熟的方案就是FP16
量化和INT8
量化,NVIDIA TensorRT
等框架就可以支持,这里不再赘述。训练方面的方案就是混合精度训练,它的基本思想很简单: 精度减半(FP32→ FP16
) ,训练时间减半。与单精度浮点数float32
(32bit,4个字节)相比,半精度浮点数float16
仅有16bit,2个字节组成。可以很明显的看到,使用FP16
可以解决或者缓解上面FP32
的两个问题:显存占用更少:通用的模型FP16
占用的内存只需原来的一半,训练的时候可以使用更大的batchsize
。计算速度更快:有论文指出半精度的计算吞吐量可以是单精度的 2-8 倍。
当前很多NVIDIA GPU
搭载了专门为快速FP16
矩阵运算设计的特殊用途Tensor Core
,比如Tesla P100
,Tesla V100
、Tesla A100
、GTX 20XX
和RTX 30XX
等。Tensor Core
是一种矩阵乘累加的计算单元,每个Tensor Core
每个时钟执行64个浮点混合精度操作(FP16
矩阵相乘和FP32
累加),英伟达宣称使用Tensor Core
进行矩阵运算可以轻易的提速,同时降低一半的显存访问和存储。随着Tensor Core
的普及FP16
计算也一步步走向成熟,低精度计算也是未来深度学习的一个重要趋势。
自动混合精度训练
不同于在推理过程中直接削减权重精度,在模型训练的过程中,直接使用半精度进行计算会导致的两个问题的处理:舍入误差(Rounding Error
)和溢出错误(Grad Overflow / Underflow
)。
舍入误差: float16
的最大舍入误差约为 (~2 ^-10
),比float32
的最大舍入误差(~2 ^-23
) 要大不少。 对足够小的浮点数执行的任何操作都会将该值四舍五入到零,在反向传播中很多甚至大多数梯度更新值都非常小,但不为零。 在反向传播中舍入误差累积可以把这些数字变成0
或者 nan
, 这会导致不准确的梯度更新,影响网络的收敛。
溢出错误: 由于float16
的有效的动态范围约为 ( 5.96×10^-8 ~ 6.55×10^4
),比单精度的float32
(1.4x10^-45 ~ 1.7x10^38
)要狭窄很多,精度下降(小数点后16相比较小数点后8位要精确的多)会导致得到的值大于或者小于fp16
的有效动态范围,也就是上溢出或者下溢出。在深度学习中,由于激活函数的的梯度往往要比权重梯度小,更易出现下溢出的情况。2018年ICLR论文
Mixed Precision Training 中提到,简单的在每个地方使用FP16
会损失掉梯度更新小于2^-24
的值——大约占他们的示例网络所有梯度更新的5%。
解决方案就是使用混合精度训练(Mixed Precision
)和损失缩放(Loss Scaling
):
-
混合精度训练: 混合精度训练是一种通过在
FP16
上执行尽可能多的操作来大幅度减少神经网络训练时间的技术,在像线性层或是卷积操作上,FP16
运算较快,但像Reduction
运算又需要FP32
的动态范围。通过混合精度训练的方式,便可以在部分运算操作使用FP16
,另一部分则使用FP32
,混合精度功能会尝试为每个运算使用相匹配的数据类型,在内存中用FP16
做储存和乘法从而加速计算,用FP32
做累加避免舍入误差。这样在权重更新的时候就不会出现舍入误差导致更新失败,混合精度训练的策略有效地缓解了舍入误差的问题。 -
损失缩放: 即使用了混合精度训练,还是会存在无法收敛的情况,原因是激活梯度的值太小,造成了下溢出。损失缩放是指在执行反向传播之前,将损失函数的输出乘以某个标量数(论文建议从8开始)。 乘性增加的损失值产生乘性增加的梯度更新值,提升许多梯度更新值到超过
FP16
的安全阈值2^-24
。 只要确保在应用梯度更新之前撤消缩放,并且不要选择一个太大的缩放以至于产生inf
权重更新(上溢出) ,从而导致网络向相反的方向发散。
使用Pytorch AMP
Pytorch
原生的amp
模式使用起来相当简单,只需要从torch.cuda.amp
导入GradScaler
和 autocast
这两个函数即可。torch.cuda.amp
的名字意味着这个功能只能在cuda
上使用,事实上,这个功能正是NVIDIA
的开发人员贡献到PyTorch
项目中的。
Pytorch
在amp
模式下维护两个权重矩阵的副本,一个主副本用 FP32
,一个半精度副本用 FP16
。 梯度更新使用FP16
矩阵计算,但更新于 FP32
矩阵。 这使得应用梯度更新更加安全。
autocast
上下文管理器实现了 FP32
到FP16
的转换,它会自动判别哪些层可以进行FP16
哪些层不可以。
GradScaler
对梯度更新计算(检查是否溢出)和优化器(将丢弃的batches转换为 no-op)进行控制,通过放大loss
的值来防止梯度的溢出。
在训练中的具体使用方法如下所示:
def train():
batch_size = 8
epochs = 10
lr = 1e-3
size = 256
num_class = 35
use_amp = True
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('torch version: {}'.format(torch.__version__))
print('amp: {}'.format(use_amp))
print('device: {}'.format(device))
print('epochs: {}'.format(epochs))
print('learn rate: {}'.format(lr))
print('batch size: {}'.format(batch_size))
net = ERFNet(num_classes=num_class).to(device)
train_data = CityScapesDataset('D:\\dataset\\cityscapes',
'D:\\dataset\\cityscapes\\trainImages.txt',
'D:\\dataset\\cityscapes\\trainLabels.txt',
size, num_class)
val_data = CityScapesDataset('D:\\dataset\\cityscapes',
'D:\\dataset\\cityscapes\\valImages.txt',
'D:\\dataset\\cityscapes\\valLabels.txt',
size, num_class)
train_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=False, num_workers=8)
val_dataloader = DataLoader(val_data, batch_size=batch_size, shuffle=False, num_workers=4)
opt = torch.optim.Adam(net.parameters(), lr=lr)
criterion = torch.nn.CrossEntropyLoss(ignore_index=255)
if use_amp:
scaler = torch.cuda.amp.GradScaler()
writer = SummaryWriter("summary")
train_loss = AverageMeter()
val_acc = AverageMeter()
val_miou = AverageMeter()
for epoch in range(0, epochs):
train_loss.reset()
val_acc.reset()
val_miou.reset()
with tqdm(total=train_data.__len__(), unit='img', desc="Epoch {}/{}".format(epoch + 1, epochs)) as pbar:
# train
net.train()
for img, mask in train_dataloader:
img = img.to(device)
mask = mask.to(device)
n = img.size()[0]
opt.zero_grad()
if use_amp:
with torch.cuda.amp.autocast():
output = net(img)
loss = criterion(output, mask)
scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()
else:
output = net(img)
loss = criterion(output, mask)
loss.backward()
opt.step()
train_loss.update(loss.item(), n)
pbar.set_postfix(**{"loss": train_loss.avg})
pbar.update(img.size()[0])
writer.add_scalar('train_loss', train_loss.avg, epoch)
# eval
net.eval()
for img, mask in val_dataloader:
img = img.to(device)
mask = mask
n = img.size()[0]
output = net(img)
pred_mask = torch.softmax(output, dim=1)
pred_mask = pred_mask.detach().cpu().numpy()
pred_mask = np.argmax(pred_mask, axis=1)
true_mask = mask.numpy()
acc, acc_cls, mean_iu, fwavacc = evaluate(pred_mask, true_mask, num_class)
val_acc.update(acc)
val_miou.update(mean_iu)
writer.add_scalar('val_acc', val_acc.avg, epoch)
writer.add_scalar('val_miou', val_miou.avg, epoch)
pbar.set_postfix(**{"loss": train_loss.avg, "val_acc": val_acc.avg, "val_miou": val_miou.avg})
实验
硬件使用NVIDIA Geforce RTX 3070
作为测试卡,这块卡有184个Tensor Core
,能比较好的支持amp
模式。模型使用ERFNet
分割模型作为基准,cityscapes
作为测试数据,10个epoch
下的测试效果如下所示:
在模型的训练性能方面,amp
模式下的平均训练时间并没有明显节省,甚至还略低于正常模式。显存的占用大约节省了25%,对于需要大量显存的模型来说这个提升还是相当可观的。理论上训练速度应该也是有提升的,我到Pytorch
的GitHub issue
里翻了一下,好像30系显卡会存在速度提不上来的问题,不太清楚是驱动支持不到位还是软件适配不到位。
Metrics | time | memory |
---|---|---|
AMP | 66.72s | 2.5G |
NO_AMP | 65.64s | 3.3G |
在模型的精度方面,我在不进行数据shuffle
的情况下统计了10个epoch
下两种模式的train_loss
和val_acc
,可以看出不管是训练还是推理,amp
模式并没有带来明显的精度损失。
网友评论