Vectorize Auto Contrast.

[sebastian-sz]

Short Description
Testing the water with this issue as #1373 is merged.

How do you feel about migrating more layers to use vectorized approach? Below is a simple performance comparison of vectorized vs current.

EDIT: misclick and submitted the issue too early...

Eager performance is better.
Graph performance is the same.

[bhack]

Can you do the same with RandomContrast? Thanks

[sebastian-sz]

@bhack Similar behaviour:

1. Eager: vectorized is quite faster
2. Graph (tf.function): similar, but this time original is a bit faster.
3. Graph (XLA): similar, but vectorized is a bit faster.
   
   
   

[bhack]

Are these on CPU or on GPU?

[sebastian-sz]

CPU

[bhack]

Is this the result of RandomContrast with a single value or a list/tuple for the factor param?

[sebastian-sz]

Single value: RandomContrast(0.5) - I don't think this should affect performance?

[bhack]

I suppose not:

keras-cv/keras_cv/layers/preprocessing/random_contrast.py

Lines 58 to 66 in 4e2d660

	if isinstance(factor, (tuple, list)):
	min = 1 - factor[0]
	max = 1 + factor[1]
	else:
	min = 1 - factor
	max = 1 + factor
	self.factor_input = factor
	self.factor = preprocessing.parse_factor((min, max), min_value=-1, max_value=2)
	self.seed = seed

So It could be interesting to test, please also try to use realistic enough input (WxH) size for production use cases.

[sebastian-sz]

@bhack I don't think there is a bug:
self.factor = preprocessing.parse_factor((min, max), min_value=-1, max_value=2) is parsed into UniformFactorSampler, which returns a float, sampled from a uniform random distribution upon call.

[sebastian-sz]

please also try to use realistic enough input (WxH) size for production use cases.

Here is a basic script that I'm using:

# layer = RandomContrast((0,1))
layer = VectorizedRandomContrast((0,1))

@tf.function
def graph_mode(inputs):
    return layer(inputs)

@tf.function(jit_compile=True)
def xla_mode(inputs):
    return layer(inputs)

for batch_size in range(8, 132, 8):
    input_shape = (batch_size, 224, 224, 3)
    batch = tf.random.uniform(input_shape)

    for _ in range(10):
        layer(batch)

    results = []
    for _ in range(100):
        batch = tf.random.uniform(input_shape)
        start = time.perf_counter()
        layer(batch)
        stop = time.perf_counter()
        results.append(stop-start)

    print(batch_size, float(tf.reduce_mean(results).numpy()))

[sebastian-sz]

@bhack Sorry, my mistake in the previous comment. I'm using tuple for factor, not 0.5.

[bhack]

Ok sorry my bad, you are right it was a call not a variable.

[bhack]

@sebastian-sz I still don't undertand with your code example:

import tensorflow as tf
from tensorflow import keras
from keras_cv.layers.preprocessing import RandomContrast

layer = RandomContrast((0,1))
input_shape = (20, 224, 224, 3)
batch = tf.random.uniform(input_shape)
layer(batch)

And if you add a debug print tf.print(contrast_factor) I still see only a single value.

keras-cv/keras_cv/layers/preprocessing/random_contrast.py

Lines 72 to 73 in b203717

	def augment_image(self, image, transformation, **kwargs):
	contrast_factor = transformation

[sebastian-sz]

@bhack Sorry for the confusion. The 0,1 argument was meant to be value_range but I passed it as a factor by accident. This is only a benchmark so I don't think the value shift matters.

In this case the 0,1 is treated as min and max values, which after slight modifications:

keras-cv/keras_cv/layers/preprocessing/random_contrast.py

Lines 59 to 61 in b203717

	if isinstance(factor, (tuple, list)):
	min = 1 - factor[0]
	max = 1 + factor[1]

Are passed to uniform factor sampler:

keras-cv/keras_cv/core/factor_sampler/uniform_factor_sampler.py

Lines 39 to 47 in 4b08a7d

	def __init__(self, lower, upper, seed=None):
	self.lower = lower
	self.upper = upper
	self.seed = seed
	def __call__(self, shape=(), dtype="float32"):
	return tf.random.uniform(
	shape, seed=self.seed, minval=self.lower, maxval=self.upper, dtype=dtype
	)

This sampler seems to be only a wrapper around tf.random.uniform to set it's min and max values.
Hence calling self.factor() is equivalent to calling tf.random.uniform with the default scalar shape ().

This is the value that can be printed via tf.print(contrast_factor) and it's the value passed to tf.image.adjust_contrast.

Currently, in case of vectorized implementation, if we use this self.factor() + tf.image.adjust_contrast (if I'm not mistaken) every image in the batch receives the same contrast_factor.

[bhack]

So where is the within the batch randomization in RandomContrast?

[sebastian-sz]

@bhack Wait, it's not called per sample?

[bhack]

I have used that gist.
I supposed the we have 1 factor for each sample in the batch but it print only a single scalar

[sebastian-sz]

From the code it looks like it is sampling new scalar per image:

keras-cv/keras_cv/layers/preprocessing/base_image_augmentation_layer.py

Lines 459 to 460 in b203717

	def _batch_augment(self, inputs):
	return self._map_fn(self._augment, inputs)

_augment calls get_random_transformation, but on my end it's also printing only a single scalar, regardless of batch size.

[bhack]

@LukeWood ?

[bhack]

@sebastian-sz I suppose that map_fn expect a stacked elem. But who is stacking factor?

elems: A tensor or (possibly nested) sequence of tensors, each of which will be unstacked along their first dimension. 

[bhack]

_augment calls get_random_transformation, but on my end it's also printing only a single scalar, regardless of batch size.

_augment is only called once for the whole batch and now I am very confused on how the map_fn call was organized.

[bhack]

@sebastian-sz Check the analysis at https://github.com/keras-team/keras-cv/issues#issuecomment-1424701740.

It is why this benchmark seems so unexpected.

[LukeWood]

@sebastian-sz feel free to contribute this! Thanks a lot for the contribution!

In general we should look to vectorize layers unless there’s a clear reason not to.