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容器时才支持自动混合精度：

您需要使用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版本中可以使用。