我正在尝试使用 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 中实现？
【问题讨论】：
标签： python tensorflow keras tf.keras
目前，仅在使用 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 变体），找不到白名单操作，我看不出性能差异。
【讨论】：