Support for Mixed Precision Training

[MahmoudAshraf97]

Hello
I implemented an autoencoder architecture for image compression and used the GDN module of the TFC library, the problem is the TFC library causes problems that prevent training with Mixed precision or half precision, when replacing GDN activation with RELU or anything else the training works perfectly with mixed precision.
Please update the code to enable mixed precision
my code for reference with colab notebook to try
https://github.com/MahmoudAshraf97/AutoencoderCompression

Thanks

[MahmoudAshraf97]

Solved by replacing "tf.convert_to_tensor" calls with "tf.cast" calls using the exact same parameters in GDN.py

[jonaballe]

Hi, can you achieve the desired functionality by passing the dtype keyword argument when instantiating the GDN layer, e.g. GDN(name="gdn_0", dtype=tf.float16)?

[MahmoudAshraf97]

this solved the GDN layer problem but caused other problems with NoisyDeepFactorized, the prior must be initialized in float16 and thus expecting all calculations to be in fp16 which is true when using fp16 precision but causes problems with data types when using mixed precision especially as tf.range function doesn't have a gpu kernel for fp16, I tried to trace the data types and fix the collisions using explicit casts but this is not a good practice as the modified library will be usable on my device only

tensorflow.python.framework.errors_impl.NotFoundError: Could not find device for node: {{node Range}} = Range[Tidx=DT_HALF]
All kernels registered for op Range:
device='CPU'; Tidx in [DT_INT64]
device='CPU'; Tidx in [DT_INT32]
device='CPU'; Tidx in [DT_DOUBLE]
device='CPU'; Tidx in [DT_FLOAT]
device='GPU'; Tidx in [DT_INT64]
device='GPU'; Tidx in [DT_INT32]
device='GPU'; Tidx in [DT_DOUBLE]
device='GPU'; Tidx in [DT_FLOAT]
[Op:Range]

please refer to #105 where I fixed the problem with GDN, DeepFactorized, and ContinuousBase, instantiating the layers in fp16 is no longer required as data types conversion is automatically handled by tensorflow and mixed precision can be enabled without changing the model code, full precision also works as intended.

Thanks