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

Question

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

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

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

Answer 1

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

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

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

您可能还需要从 Python 脚本中设置特定的environment variable ，即：

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

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

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

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

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

此外，您还可以检查来自 Google 的bfloat16 （脑浮点）类型，它具有float32 （ 8位）的exponent部分和较小的分数。 与float16相比，这允许它保持更大范围的值（例如，在计算梯度时），这可以避免loss scaling 。

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