python - 是否可以使用 float16 使用 tensorflow 1 进行训练？

Question

目前使用默认设置 - float32 在 tensorflow 模型上训练 keras。

训练后网络被量化：将权重转换为 float16。这将性能提高了 ~x3，同时保持了相同的精度。

我试图从一开始就使用 float16 进行训练，结果惨遭失败。我找不到任何链接来解释这是否可能，如果不是，为什么不可能。

Answer 1

来自 NVidia 的自动混合精度可能是一种方法。

从我收集的内容来看，1.14它在上游得到了（曾经）支持。您所要做的就是像这样包装您的优化器：

opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)

您可能还需要environment variable从 Python 脚本中进行设置，即：

os.environ[‘TF_ENABLE_AUTO_MIXED_PRECISION’] = ‘1’

以上应该已经采用了良好的混合精度训练实践（例如损失缩放、float32必要时保持等）。

这个解决方案的好资源应该是NVidia 的官方文档。

收集的其他一些资源也可能有用（尽管似乎并不表明您需要做更多的事情）here、here或here。

我建议不要手动转换，因为BatchNorm除非您知道特定层的来龙去脉，否则您可能很容易失去精度（例如在推理过程中使用的统计数据）。

此外，您还可以检查bfloat16来自 Google 的 (brain float) 类型，该类型具有(位)exponent的一部分和较小的分数。与允许避免.float328float16loss scaling

上面的 ( bfloat16) 应该主要在 TPU 中有用，AFAIK NVidia GPU 对它的支持不是太好（如果我错了，请纠正我）。这里有一些信息。