python - Tensorflow Beta 2.0 中的 tf.keras 是否支持自动混合精度？

Question

我正在尝试使用 tf.keras API 让 Tensorflow 的自动混合精度工作(在 RTX 2080 Ti 上使用张量核心)，但我看不到训练有任何加速。

我刚刚添加

os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'

到 Python 脚本的顶部。我还尝试从命令行将环境变量设置为 1，即

export TF_ENABLE_AUTO_MIXED_PRECISION=1

在这种情况下是否支持 AMP，或者模型是否需要在“原始”Tensorflow 中实现？

Answer 1

目前，只有在使用 NVIDIA 的 Tensorflow Docker 容器时才支持自动混合精度:

https://ngc.nvidia.com/catalog/containers/nvidia:tensorflow

https://www.tensorflow.org/install/docker

需要使用Ubuntu 18.04，目前不支持Ubuntu版本和Windows。如果我没记错的话，最后一个 Docker 容器有 TF 1.13。安装后，tf.keras 应该支持自动混合精度。

编辑:

我在 Windows 上试用了 2.0.0-beta1，在使用自动混合精度时也没有发现任何速度提升。使用 Linux 上的 NVIDIA Docker 容器，将 TF_ENABLE_AUTO_MIXED_PRECISION 设置为 1 时，我至少获得了 2 倍的加速。希望这将在 2.0 版本中起作用。

编辑 2:使用 TF 2.0.0-rc0，AMP 提高了简单模型的预期性能。对于更复杂的模型(U-Net 变体)，没有发现白名单操作，我看不出性能差异。