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

Question

I am trying to get Tensorflow's automatic mixed precision working (to use the tensor cores on an RTX 2080 Ti), using the tf.keras API, but I can't see any speed-up in training.

I have just added

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

to the top of the Python script. I also tried setting the environment variable to 1 from the command line, ie

export TF_ENABLE_AUTO_MIXED_PRECISION=1

Is AMP supported in this case, or does the model need to be implemented in 'raw' Tensorflow?

Answer 1

At the moment, automatic mixed precision is only supported when using the Tensorflow Docker container from NVIDIA:

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

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

You need to use Ubuntu 18.04, the current Ubuntu version and Windows are not supported. The last Docker container has TF 1.13 if I'm not mistaken. Once installed, tf.keras should support automatic mixed precision.

Edit:

I tried the 2.0.0-beta1 on Windows and also did not notice any speed improvement when using automatic mixed precision. With the NVIDIA Docker container on Linux, I got at least 2x speedup when setting TF_ENABLE_AUTO_MIXED_PRECISION to 1. Hopefully, this will work in the 2.0 release.

Edit 2: With TF 2.0.0-rc0, AMP improves the performance as expected for a simple model. With a more complex model (a U-Net variant), no whitelist ops are found and I see no performance difference.