Add vectorized RandomBrightness, RandomContrast, RandomHue and RandomColorJitter

[james77777778]

What does this PR do?

Fixes #1405
Fixes #1376

Following #1392 but I moved A/B testing into benchmarks suggested by @ianstenbit 's comment at #1392 (comment), also I added XLA Mode jit_compile=True for benchmarks suggested by @bhack

For unit tests, only keras_cv/layers/preprocessing/random_hue_test.py failed at test_adjust_full_opposite_hue and could be resolved by relaxing constraints atol=1e-5, rtol=1e-5

For numerical check (used image_shape = (1024, 32, 32, 3) locally):

• ./benchmarks/vectorized_random_brightness.py passes without any constraint relaxing
• ./benchmarks/vectorized_random_contrast.py passes with atol=1e-5, rtol=1e-5 at rgb value range
• ./benchmarks/vectorized_random_hue.py passes with atol=1e-3, rtol=1e-5 at rgb value range

./benchmarks/vectorized_random_color_jitter.py has no numerical check because it is composited by above layers

Some notes:

• OldRandomBrightness got similar performance with Graph/XLA Mode, other vectorized layers got significant boost
• RandomHue failed to run XLA mode
• The benchmark of RandomColorJitter was done after I vectorizing all sub-layers

This is results from benchmarks:

RandomBrightness

RandomContrast

RandomHue (no XLA Mode)

RandomColorJitter

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue? Please add a link
  to it if that's the case.
• Did you write any new necessary tests?
• If this adds a new model, can you run a few training steps on TPU in Colab to ensure that no XLA incompatible OP are used?

Who can review?

@LukeWood @ianstenbit

[ianstenbit]

@james77777778 this is great -- thank you!

A few quick questions:

For RandomHue -- did XLA also fail for the original implementation? Or only for the new vectorized implementation?

For RandomBrightness -- the graph seems to be missing a few of the trials. I assume this is because the scale of the Y axis is too small? Would you mind updating that?

Thank you for this great contribution!

[ianstenbit]

/gcbrun

[james77777778]

@ianstenbit

For RandomHue -- did XLA also fail for the original implementation? Or only for the new vectorized implementation?

• RandomHue (vectorized): ran w/o problem
• OldRandomHue: failed for XLA

Code snippet to bypass failure

        # XLA Mode
        if aug is OldRandomHue:
            print("Skip XLA Mode")
            continue
        c = aug.__name__ + " XLA Mode"
        layer = aug(**aug_args)

For RandomBrightness -- the graph seems to be missing a few of the trials. I assume this is because the scale of the Y axis is too small? Would you mind updating that?

Yes, the Y axis is too small to distinguish the differences.

From logs:

Mode	N	RandomBrightness	OldRandomBrightness
Eager	1000	0.0117	0.0818
	2000	0.0135	0.0866
	5000	0.0611	0.1357
	10000	0.1203	0.1954
Graph	1000	0.0105	0.0072
	2000	0.0116	0.0134
	5000	0.0604	0.0604
	10000	0.1199	0.1201
XLA	1000	0.0072	0.0073
	2000	0.0117	0.0137
	5000	0.0608	0.0609
	10000	0.1201	0.1213

The vectorized RandomBrightness won in Eager mode but got similar performance in Graph/XLA mode.

Thanks for the review!