我正在尝试使用 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 中实现?
请您参考如下方法:
目前,只有在使用 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 变体),没有发现白名单操作,我看不出性能差异。
