Mixed precision and XLA

[dathudeptrai]

It's indispensable that we could make all models be able to take advantage of mixed_precision and xla. In my experience, mixed_precision and xla can speed up training and inference time around two times. Also, I think if you want this repo to become more popular, you should show users how to maximize the performance of TensorFlow. If the training time and inference time of all models here are not competitive with other frameworks then there is no reason for users to use this framework. From what I can see, 95% of TensorFlow users don't know how to maximize the performance of TF.

@LukeWood @ianstenbit

[bhack]

Some of my points here (vectorized_map + XLA):
#581 (comment)

We could add mixed_precision tests decorator/coverage as an extra point.

[ianstenbit]

Hey @dathudeptrai -- thanks for bringing this up.

I agree that performance improvements of our models and strong example content to demonstrate performance optimization are important for the usability of KerasCV.

Our example training scripts currently do use XLA by default, and adding mixed_precision support would be a very welcome contribution if you're interested.

If you encounter any models which do not support either XLA or mixed precision out of the box, please feel free to create an issue or PR addressing it.

[bhack]

Our example training scripts currently do use XLA by default, and adding mixed_precision support would be a very welcome contribution if you're interested.

@ianstenbit I think this is not the case as this is just going to jit_compile the train_step right?

E.g. if you are going to test under jit_compile the layers you are using in your first/basic contributed training example it will go to fail:

https://github.com/keras-team/keras-cv/blob/master/examples/training/classification/imagenet/basic_training.py

Just to make an example for CutMix that you use in the example:
on XLA_CPU_JIT: RandomGamma (No registered 'RandomGamma' OpKernel for XLA_CPU_JIT devices compatible with node

As we don't have a RandomGamma kernel (the ops we are using in CutMix) in the TF2XLA bridge, the XLA compilation is failing.

Reminder
XLA has a fail fast approach -> the first not covered ops will let to fail the whole compilation)

/cc @cheshire @andrew-leaver

See also https://github.com/keras-team/keras/issues/16403. I suggested that @qlzh727 could migrate it on this repo as the base class is here now.

[DavidLandup0]

Agreed that most TF users don't really maximize the performance. Using jit_compile, mixed precision training, etc. are one-liners. It somehow looks like these are lowkey unknown. Using tf.data instead of ImageDataGenerator and/or custom generators, is common too but less so. It's not only about making KerasCV compatible with performance boosters, but also about raising awareness for the majority of TF users. :)

One way to do so is to add "Suggestion" messages while compiling models, but some might see that as intrusive perhaps?