训练提速60%！只需5行代码，PyTorch 1.6即将原生支持自动混合精度训练

with torch.cuda.amp.autocast():     y_pred = model(X_batch).squeeze()     loss = self.loss_fn(y_pred, y_batch) 

以这种方式包装前向传播，可以自动打开后传(如 loss.backwards ())的autocasting，因此不需要调用两次autocast。

只要你遵循PyTorch 的最佳实践(例如，避免in-place操作) ，autocasting基本上就可以“正常工作”。它甚至可以使用多GPU DistributedDataParallel API (只要遵循建议的策略，每个 GPU 只使用一个进程)。只需一个小调整，多GPU DataParallel API也可以用。Pytorch 文档中的 Automatic Mixed Precision Examples 页面的“Working with multiple GPUs”部分是关于这个主题的一个方便的参考。个人观点，有一个要记住的重点是: "优先用 binary cross entropy with logits 而不是 binary cross entropy"。

Benchmarks性能

此时，我们已经了解了什么是混合精度，什么是张量核，以及 PyTorch API 如何实现自动混合精度。唯一剩下的就是看看一些真实世界的性能benchmarks！

我曾经用自动混合精度训练过三个非常不一样的神经网络，还有一次没用，通过 Spell API 调用 V100s (上一代张量核)和 T4s (当代张量核)。我分别使用了 AWS EC2实例、 p3.2xlarge 和 g4dn.xlarge，最近的 PyTorch 1.6 nightly 和 CUDA 10.0。所有模型的收敛都是一致的，即没有一个模型发现混合精度网络和原网络在训练损失上有任何差异。训练的网络如下:

前馈, 一个前馈神经网络，训练数据来自Kaggle比赛Rossman Store Samples UNet, 一个中等大小的原版UNet 图像分割网络， 在数据集Segmented Bob Ross Images 上训练 BERT, 一个大的 NLP transformer 模型，使用bert-base-uncased 骨干(通过 huggingface)，及数据来自Kaggle竞赛 Twitter Sentiment Extraction

结果如下:

由于前馈网络非常小，混合精度训练对它没有任何好处。

UNet 是一个中等规模的卷积模型，共有7,703,497个参数，从混合精度训练中得到了显著的好处。有趣的是，虽然 V100和 T4都受益于混合精度训练，但 T4的好处要大得多: 节省5%时间vs. 高达30%的时间。

BERT 是一个很大的模型，在这里使用混合精度训练节省时间，从中等模型的“很好”到了“必须拥有”。在Volta或Turing GPU 上训练，自动混合精度将为大型模型减少50% 到60% 的训练时间！

这是一个巨大的优势，尤其是当你考虑到增加的复杂性极小时——只需要对模型训练脚本进行四到五行代码修改。在我看来:

混合精度应该是你对模型训练脚本进行的最先性能优化之一。

内存呢？

正如我在“混合精度是如何工作的”一节中解释的那样，在内存中fp16矩阵的大小是fp32矩阵的一半，因此，混合精度训练的另一个据称的优势是内存使用率。GPU 内存的瓶颈远小于 GPU 的计算能力，但仍有很大的优化价值。你的内存使用效率越高，你可以在 GPU 上使用的batch size就越大。

PyTorch 在模型训练过程开始时保留一定数量的 GPU 内存，并在训练期间保留这些内存。这可以防止其它进程在训练过程中抢占过多的 GPU 内存，迫使 PyTorch 训练脚本崩溃并出现 OOM 错误。

以下是启用混合精度训练对 PyTorch 内存保留行为的影响:

有趣的是，虽然两个较大的模型都看到了切换到混合精度的好处，UNet 从切换中得到的好处比 BERT 多得多。PyTorch 内存分配行为对我来说非常不透明，所以我不知道为什么会出现这种情况。

总结

在即将发布的 PyTorch 1.6版本中，自动混合精度训练是一个易于使用且功能强大的新特性，该版本承诺将在最新的 NVIDIA GPU 上运行的大型模型训练工作加快60% 。

虽然这种技术已经存在了一段时间，但是对于普通用户来说还不是很容易理解，因为直到现在它还没有一个原生 PyTorch API。

要直接从源代码中了解更多关于混合精度训练的信息，请参阅 PyTorch master 文档中的automatic mixed precision package和automatic mixed precision examples页面。

想自己测试一下这个功能？安装最新的 PyTorch nightly非常简单: 查看 PyTorch 主页上的说明了解如何安装。

想要自己复现这些benchmarks吗？所有模型源代码都可以在 GitHub 上的 ResidentMario/spell-feedforward-rossman, ResidentMario/spell-unet-bob-ross, 和 ResidentMario/spell-tweet-sentiment-extraction 库中获得。

（编辑：南平站长网）

【声明】本站内容均来自网络，其相关言论仅代表作者个人观点，不代表本站立场。若无意侵犯到您的权利，请及时与联系站长删除相关内容!