Enable training/inference with half-precision (float16) about deepcell-tf HOT 5 CLOSED

As discussed in #436, there are float16 issues with both the ImageDataGenerators and the model zoo.

from deepcell-tf.

It looks like this is supported as mixed precision in TensorFlow 2.

from deepcell-tf.

This is indeed supported natively by TensorFlow 2:

from tensorflow.keras import mixed_precision

# On TPUs, use 'mixed_bfloat16' instead
mixed_precision.experimental.set_policy('mixed_float16')

# On TensorFlow 2.4.1:
# mixed_precision.set_global_policy('mixed_float16')

I tried training the Nuclear Segmentation notebook using the mixed-precision above. However, I found that it caused a significant slowdown during training. (epoch goes from 10 minutes to 60+ minutes). I'm not sure if this is due to our data generators or something about our PanOpticNet model architecture - more debugging is required.

from deepcell-tf.

Mixed precision works in TensorFlow 2.4.1!

However, it looks like std and max normalizations have an issue with mixed precision:

TypeError: Input 'filter' of 'DepthwiseConv2dNative' Op has type float32 that does not match type float16 of argument 'input'.

The root cause is that the layer's kernel was initialized using K.floatx which is not changed when setting the global precision policy. This is resolved in #490

from deepcell-tf.

While #490 does allow mixed-precision training without any warnings or errors, it does not seem to improve training performance. I get roughly the same training time (running the Nuclear Segmentation PanOpticNet notebook) using mixed precision and float32 while testing different settings (Location, normalization method, interpolation upsample_type, etc).

That said, mixed-precision DOES enable much larger batch sizes than with float32 mode, so it is indeed working.

from deepcell-tf.