Skip to main content
 首页 » 编程设计

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

2024年11月24日26telwanggs

我正在尝试使用 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 变体),没有发现白名单操作,我看不出性能差异。