使用 TensorFlow 混合精度训练的首次实验

上周发布的 TensorFlow 2.1 允许进行混合精度训练,利用最新的 NVidia GPU 中提供的 Tensor Cores。在这篇文章中,我们报告了第一批实验结果,并提供了有关这一切的背景信息。

来源:RStudio AI博客

从最近的 2.1 版本开始,TensorFlow 支持 Keras 的混合精度训练(以下简称 MPT)。在这篇文章中,我们尝试了 MPT 并提供了一些背景信息。首先说明:在 Tesla V100 GPU 上,我们基于 CNN 的实验并未显示执行时间有显著减少。在这种情况下,很难决定是否真的要写一篇文章。您可以争辩说,就像在科学中一样,零结果也是结果。或者,更实际地说:它们开启了一场讨论,可能导致发现错误、澄清使用说明和进一步实验等。

混合精度训练 null

此外,这个话题本身就足够有趣,值得进行一些背景解释——即使结果还没有完全出现。

尚未

首先,让我们来听听 MPT 的一些背景信息。

这不仅仅是为了节省内存

在 TensorFlow 中描述 MPT 的一种方式可能是这样的:MPT 允许您训练权重类型为 float32 或 float64 的模型,这通常(出于数值稳定性的原因),但数据(操作之间推送的张量)的精度较低,即 16 位(float16)。

权重 float32 float64 数据 float16

这句话可能适合作为 TLDR;对于较新的 MPT 文档页面,也可以在 TensorFlow for R 网站上找到 R 文档。根据这句话,您可能会想“哦,当然,这是为了节省内存”。更少的内存使用量意味着您可以运行更大的批量大小而不会出现内存不足错误。

TLDR; MPT 文档页面 TensorFlow for R 网站 “哦,当然,所以这是为了节省内存”

这当然是正确的,您会在实验结果中看到它。但这只是故事的一部分。另一部分与 GPU 架构和并行(不仅仅是 GPU 上的并行,正如我们将看到的)计算有关。

AVX & co.

SIMD(单指令多数据)
# 将这些视为 64 位整数 c(1, 2) + c(3, 4)
# 将这些视为 64 位整数 c