Augment benchmarks/vectorization_strategy_benchmark.py to include XLA compilation

[LukeWood]

No description provided.

[bhack]

Check:
https://github.com/keras-team/keras-cv/pull/161/files#r820872659

[bhack]

It seems that in the faster jit_compiled version all the tf.cast we have are still quite relevant

[bhack]

@qlzh727 What do you think about adding an extra parameter to the base class for jit_compile?
https://github.com/keras-team/keras/blob/master/keras/layers/preprocessing/image_preprocessing.py#L413-L414

So that we could optionally use

f = def_function.function(self._augment, jit_compile=True)
self._map_fn(f, inputs)

[LukeWood]

@qlzh727 What do you think about adding an extra parameter to the base class for jit_compile? https://github.com/keras-team/keras/blob/master/keras/layers/preprocessing/image_preprocessing.py#L413-L414

So that we could optionally use

f = def_function.function(self._augment, jit_compile=True)
self._map_fn(f, inputs)

I personally would support a jit_compile attribute on the BaseImageAugmentationLayer.

so in a constructor you'd do:

jit_compile = True

and it would work pass it to the tf.function wrapping call.

[bhack]

We need to care about massively closing bugs a we lose interesting threads. This is now at keras-team/keras-io#1541

[LukeWood]

We need to care about massively closing bugs a we lose interesting threads. This is now at keras-team/keras-io#1541

This is a distinct issue. I do not think benchmarks do not need this as the benchmarks existed to guide our direction.

you are right, BaseImageAugmentationLayer should support jit_compile, but the benchmarks do not need this so I closed this issue.

Closing bug is part of maintaining a clean project state.

[bhack]

Closing bug is part of maintaining a clean project state.

Yes but if there are still interesting parts emerged in the thread you need to spinoff that part on a new ticket or it will be lost forever.