PyTorch1.6nightly添加一个子模块amp以支持自动混合精度训练。值得期待。再来看看性能,相比NvidiaApex有哪些优势?PyTorch1.6中即将推出的torch.cuda.amp混合精度训练模块实现了其承诺,只需几行新代码即可将大型模型的训练速度提高50-60%。预计PyTorch1.6中最令人兴奋的新增功能之一是支持自动混合精度训练。混合精度训练是一种通过对半精度浮点fp16执行尽可能多的操作来大幅减少神经网络训练时间的技术,它取代了PyTorch默认的单精度浮点fp32。最新一代的NVIDIAGPU具有专为快速fp16矩阵运算而设计的专用张量内核。然而,直到现在,这些TensorCore一直难以使用,因为它需要手动将降低精度的操作写入模型。这就是自动化混合精度训练的用武之地。即将推出的torch.cuda.ampAPI将允许您仅用五行代码在训练脚本中实施混合精度训练!混合精度如何工作在了解混合精度训练如何工作之前,我们首先需要回顾一下浮点数。在计算机工程中,像1.0151或566132.8这样的十进制数传统上表示为浮点数。由于我们可以拥有无??限精度的数字(想象一下π),但存储它们的空间是有限的,因此我们必须在精度(四舍五入之前我们可以包含在数字中的小数位数)和大小(我们在用于存储数字的位数之间进行了折衷)。浮点数技术标准IEEE754制定了以下标准:fp64,又名双精度或“双”,最大舍入误差~2^-52fp32,又名单精度或“单”,最大舍入误差~2^-23fp16,又名半精度或“一半”,最大舍入误差为~2^-10。Python浮点类型是fp64,而对内存更敏感的PyTorch使用fp32作为默认数据类型。混合精度训练的基本思路很简单:精度减半(fp32→fp16),训练时间减半。最难的部分是如何安全地执行此操作。请注意,较小的浮点数会导致较大的舍入误差。对“足够小”的浮点数执行的任何操作都会将该值舍入为零!这称为下溢,这是一个问题,因为反向传播期间的许多(如果不是大多数)梯度更新非常小,但不是零。反向传播过程中的舍入误差累积可以将这些数字变成0或nans;这会导致不准确的梯度更新并影响网络的收敛。2018年ICLR论文混合精度训练发现,简单地在任何地方使用fp16都会“吞噬”小于2^-24的梯度更新——大约占他们示例网络所有梯度更新的5%:混合精度训练是一组技术,可让您使用fp16而不会导致模型训练发散。这是三种不同技术的结合。首先,维护权重矩阵的两个副本,一个fp32的“主副本”和一个fp16的半精度副本。梯度更新是使用fp16矩阵计算的,但在fp32矩阵上更新。这使得应用梯度更新更加安全。其次,不同的向量运算以不同的速率累积错误,所以区别对待。有些操作在fp16中始终是安全的,而其他操作仅在fp32中是安全的。不是使用fp16运行整个神经网络,有些使用半精度,有些使用单精度。这种数据类型的混合就是为什么这种技术被称为“混合精度”的原因。第三,使用损失缩放。损失缩放是指在执行反向传播之前,将损失函数的输出乘以某个标量数(论文建议从8开始)。乘法增加损失值会产生乘法增加的梯度更新值,“提升”许多梯度更新值超过2^-24的fp16安全阈值。只需确保在应用梯度更新之前撤消缩放,并且不要选择太大的缩放以产生inf权重更新(溢出),从而导致网络向相反的方向发散。将这三种技术结合在一起,作者可以训练多种网络在显着加速的时间内收敛。至于基准,我建议阅读这篇9页的论文!TensorCores的工作原理虽然混合精度训练可以节省内存(fp16矩阵的大小是fp32矩阵的一半),但如果没有特殊的GPU支持,它不会加速模型训练。芯片上需要有可以加速半精度运算的东西。在最近几代的NVIDIAGPU中,这个东西被称为:Tensorcores。张量核心是一种新型处理单元,针对非常具体的操作进行了优化:将两个4×4fp16矩阵相乘,并将结果添加到第三个4×4fp16或fp32矩阵(一种“融合乘加(fusedmultiplyadd)”).可以使用此运算作为其基本构建块来实现更大的fp16矩阵乘法运算。由于大多数反向传播归结为矩阵乘法,因此TensorCores适用于网络中几乎任何计算密集型层。陷阱:输入矩阵必须是fp16。如果您在具有TensorCores的GPU上进行训练而不是混合精度训练,您将无法从显卡获得100%的回报!fp32中定义的标准PyTorch模型永远不会将任何fp16数学应用于芯片,因此所有那些极其强大的张量核心都将闲置。Tensor核心于2017年底在上一代Volta架构中引入,当前一代Turing已经看到一些改进,并将在即将推出的Ampere中看到进一步改进。云端普遍可用的两款GPU分别是V100(5120个CUDA核心,600个张量核心)和T4(2560个CUDA核心,320个张量核心)。另一个需要记住的难题是固件。虽然CUDA7.0或更高版本支持TensorCore操作,但据说较早的实现存在bug,因此使用CUDA10.0或更高版本很重要。有了关于Pytorch的自动混合精度如何工作的重要背景,我们终于可以开始深入研究新的PyTorchampAPI。混合精度训练在技术上永远是可能的:在fp16中手动运行部分网络,并自己实现损失缩放。自动混合精度训练令人兴奋的部分是“自动”部分。只需学习几个新的API原语:torch.cuda.amp.GradScalar和torch.cuda.amp.autocast。启用混合精度训练就像在训练脚本中插入正确的位置一样简单!为了演示,下面是一段用于混合精度训练的网络训练循环的代码。#NEW标记标记添加了新代码的位置。self.train()X=torch.tensor(X,dtype=torch.float32)y=torch.tensor(y,dtype=torch.float32)优化器=torch.optim.Adam(self.parameters(),lr=self.max_lr)scheduler=torch.optim.lr_scheduler.OneCycleLR(optimizer,self.max_lr,cycle_momentum=False,epochs=self.n_epochs,steps_per_epoch=int(np.ceil(len(X)/self.batch_size)),)批次=torch.utils.data.DataLoader(torch.utils.data.TensorDataset(X,y),batch_size=self.batch_size,shuffle=True)#NEWscaler=torch.cuda.amp.GradScaler()forepochinrange(self.n_epochs):fori,(X_batch,y_batch)inenumerate(batches):X_batch=X_batch.cuda()y_batch=y_batch.cuda()optimizer.zero_grad()#NEWwithtorch.cuda.amp.autocast():y_pred=模型(X_batch)。squeeze()loss=self.loss_fn(y_pred,y_batch)#NEWscaler.scale(loss).backward()lv=loss.detach().cpu().numpy()ifi%100==0:print(f"Epoch{epoch+1}/{self.n_epochs};Batch{i};Loss{lv}")#NEWscaler.step(optimizer)scaler.update()scheduler.s新的PyTorchGradScaler对象是PyTorch的破坏缩小。回想一下“混合精度如何工作”部分,在训练期间需要某种形式的缩放以防止梯度收缩为零。最佳损失乘数足够高以保留非常小的梯度,但不会高到导致非常大的梯度四舍五入到inf从而导致相反的问题。PyTorch使用指数退避算法来解决这个问题。Gradscalar从一个每次加倍的小损失乘数开始。这种逐渐加倍的行为一直持续到GradScalar遇到包含inf值的梯度更新。Gradscalar丢弃这批数据(例如跳过梯度更新),将损失乘数减半,并重置其倍增时间。通过以这种方式上下移动损失乘数,PyTorch可以随时间逼近适当的损失乘数。熟悉TCP拥塞控制的读者应该对这里的核心思想非常熟悉!算法使用的确切数量是可配置的,您可以直接从文档字符串中看到默认值:torch.cuda.amp.GradScaler(init_scale=65536.0,growth_factor=2.0,backoff_factor=0.5,growth_interval=2000,enabled=True)Gradscalar需要控制梯度更新计算(检查溢出)和优化器(将丢弃的批次转换为无操作)以进行操作。这就是为什么loss.backwards()被scaler.scale(loss).backwards()代替,optimizer.step()被scaler.step(optimizer)代替的原因。值得注意的是,GradScalar可以检测和阻止溢出(因为inf总是坏的),但它不能检测和阻止下溢(因为0通常是合法值)。如果您选择的初始值太低且增长间隔太长,您的网络可能会在GradScalar介入之前发生下溢和发散。因此,选择一个非常大的初始值可能是个好主意。最后,请注意GradScalar是一个有状态对象。使用此功能保存模型检查点需要写入和读取磁盘以及模型权重。这可以使用state_dict和load_state_dict对象方法轻松完成(在PyTorch文档中描述)。自动混合精度训练难题的另一半是torch.cuda.amp.autocast上下文管理器。Autocast实现了fp32->fp16转换。回忆一下“混合精度的工作原理”,并非所有操作都可以在fp16中安全运行,因为不同的操作以不同的速率累积错误。下面的屏幕截图来自amp模块文档,显示了autocast如何处理PyTorch中可用的各种操作:该列表主要由矩阵乘法和卷积以及简单的线性函数组成。这些操作在fp16中是安全的,但有向上转换规则以防输入混合使用fp16和fp32以确保它们不会中断。请注意,此列表还包括其他两个基本线性代数运算:矩阵/向量点积和向量叉积。对数、指数、三角函数、正态函数、离散函数和(大)和在fp16中不安全,必须在fp32中实现。通过这个列表,在我看来,由于它们内部依赖于基本的线性代数运算,大多数层都会从自动转换中受益,但大多数激活函数却没有。卷积层是大赢家。启用sutocasting非常简单。您需要做的就是使用autocast上下文管理器来包装模型的前向传播:withtorch.cuda.amp.autocast():y_pred=model(X_batch).squeeze()loss=self.loss_fn(y_pred,y_batch)to这种包裹forwardpass的方式可以自动开启backpass的autocasting(比如loss??.backwards()),所以不需要两次调用autocast。只要您遵循PyTorch最佳实践(例如,避免就地操作),自动投射基本上“就可以工作”。它甚至可以使用多GPUDistributedDataParallelAPI(只要您遵循每个GPU仅使用一个进程的推荐策略)。通过一个小的调整,多GPUDataParallelAPI也可用。Pytorch文档中自动混合精度示例页面的“使用多个GPU”部分是关于此主题的方便参考。从我个人的角度来看,有一点需要记住:“更喜欢使用带有logits的二元交叉熵而不是二元交叉熵”。基准性能至此,我们了解了什么是混合精度,什么是张量核,以及PyTorchAPI如何实现自动混合精度。唯一剩下的就是查看一些真实世界的性能基准!我用自动混合精度训练了三种截然不同的神经网络,还有一种不起作用,通过SpellAPI调用V100s(上一代张量核心)和T4s(当前一代张量核心)。我分别使用了AWSEC2实例、p3.2xlarge和g4dn.xlarge、最近的PyTorch1.6nightly和CUDA10.0。所有模型的收敛性都是一致的,即没有一个模型发现混合精度网络和原始网络之间的训练损失有任何差异。训练好的网络如下:Feedforward,前馈神经网络,训练数据来自Kaggle竞赛RossmanStoreSamplesUNet,中型原始UNet图像分割网络,BERT在数据集SegmentedBobRossImages,大型NLP上训练transformer模型,使用的是bert-base-uncasedbackbone(viahuggingface),来自Kaggle比赛的数据TwitterSentimentExtraction结果如下:由于前馈网络很小,mixed-precisiontrainingdoesnotdoitany好的。UNet是一个中等大小的卷积模型,共有7,703,497个参数,从混合精度训练中受益匪浅。有趣的是,虽然V100和T4都受益于混合精度训练,但T4的好处要大得多:节省5%的时间对比高达30%的时间。BERT是一个大模型,这里使用混合精度训练可以节省时间,对于平庸的模型从“不错”到“必备”。在Volta或TuringGPU上训练,自动混合精度将使大型模型的训练时间减少50%到60%!这是一个巨大的优势,尤其是当您认为增加的复杂性很小时——只需要对模型训练脚本修改四五行代码。在我看来:混合精度应该是您对模型训练脚本进行的首要性能优化之一。内存呢?正如我在“混合精度如何工作”部分中解释的那样,fp16矩阵在内存中的大小是fp32矩阵的一半,因此混合精度训练的另一个据称优势是内存使用。GPU显存的瓶颈远小于GPU的计算能力,但仍有很大的优化价值。您的内存使用效率越高,您可以在GPU上使用的批处理大小就越大。PyTorch在模型训练过程开始时预留一定量的GPU内存,并在训练过程中预留这部分内存。这可以防止其他进程在训练期间占用过多的GPU内存,从而迫使PyTorch训练脚本因OOM错误而崩溃。以下是启用混合精度训练对PyTorch的内存保留行为的影响:有趣的是,虽然两个较大的模型都看到了切换到混合精度的好处,但UNet从切换中获得的好处明显多于BERT。PyTorch内存分配行为对我来说非常不透明,所以我不知道为什么会这样。摘要自动混合精度训练是即将发布的PyTorch1.6版本中一项易于使用且功能强大的新功能,它有望在最新的NVIDIAGPU上将大型模型训练作业的速度提高60%。虽然这种技术已经存在了一段时间,但它对普通用户来说并不是很容易理解,因为直到现在它还没有原生的PyTorchAPI。要直接从源代码了解有关混合精度训练的更多信息,请参阅PyTorch主文档中的自动混合精度包和自动混合精度示例页面。想亲自测试一下这个功能吗?每晚安装最新的PyTorch很容易:查看PyTorch主页上的说明以了解操作方法。想自己重现这些基准?所有模型源代码均可在GitHub的ResidentMario/spell-feedforward-rossman、ResidentMario/spell-unet-bob-ross和ResidentMario/spell-tweet-sentiment-extraction存储库中找到。
