Introduces MaybeApply layer. #435

sebastian-sz · 2022-05-15T09:15:22Z

Maybe Apply layer is a wrapper around native Keras layers or BaseImageAugmentationLayer that applies, specified operation to random samples in a batch.

Example usage

images = tf.random.stateless_uniform(shape=(5, 2, 2, 1), seed=[0, 1])

zero_out = tf.keras.layers.Lambda(lambda x: 0 * x)
maybe_apply = MaybeApply(layer=zero_out, rate=0.5, seed=1234)

outputs = maybe_apply(images)
print(outputs)  # Random samples in a batch have been zero'ed out.

Performance / overhead

The layer introduces an overhead for the layer it wraps. It looks like the layer wrapped in MaybeApply takes ~2x the time to execute. I'm not sure if any improvements here can be made - I tried to implement the same behaviour using tf.gather_nd + tf.scatter_nd_update but got similar or even worse results for larger batches.

Below are latency (ms) measurements for Solarization and Posterization. Both were benchmarked with XLA to avoid tf.function retracing.

Known Issues

The layer throws error in XLA with auto_vectorize=True, regardless of what layer it wraps. I'm not really sure why

InvalidArgumentError: Reading input as constant from a dynamic tensor is not yet supported. Xla shape: s32[<=32]

but it works with auto_vectorize=False.

bhack · 2022-05-15T09:55:09Z

It looks like the layer wrapped in MaybeApply takes ~2x the time to execute

You could try to trace/profile the code to collect more insight:
https://www.tensorflow.org/api_docs/python/tf/profiler/experimental/Trace

sebastian-sz · 2022-05-15T11:32:10Z

@bhack
Running with XLA for both MaybeApply and raw layer I'm getting:
Your program is NOT input-bound because only 0.0% of the total step time sampled is waiting for input.
And it says that almost entire time is spent on Host Compute Time. Still this is e.g. 12.2 ms for native layer and 25.3 ms for MaybeApply wrapped.

However, running without XLA: I'm getting a different message for MaybeApply:
Your program is POTENTIALLY input-bound because 36.5% of the total step time sampled is spent on 'All Others' time (which could be due to I/O or Python execution or both).

I am not reading any data from the disk, so I'd assume it could be Python's overhead? Not sure.

bhack · 2022-05-15T12:14:09Z

Do you have a gist with tf.profile to reproduce this?

sebastian-sz · 2022-05-15T12:40:37Z

@bhack this is what I'm using:
benchmark_maybe_apply.zip

bhack · 2022-05-15T14:53:35Z

Known Issues
The layer throws error in XLA with auto_vectorize=True, regardless of what layer it wraps. I'm not really sure why
InvalidArgumentError: Reading input as constant from a dynamic tensor is not yet supported. Xla shape: s32[<=32]

Yes it Is:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/tf2xla/xla_op_kernel.cc#L138-L146

Do you have, can you post, Tensorboard screenshots for the ops section for plain vs maybe wrapper?

The ops tables are like these:
#165 (comment)

More in general with XLA is also interesting to check how well the graph ops that we have used in the implementaton are optimized-fused or not with the HLO dump (not all the possible ops permutation sequences have the same XLA coverage and optimizzation quality).
See:
#141 (comment)

bhack · 2022-05-15T20:38:38Z

@bhack this is what I'm using:
benchmark_maybe_apply.zip

You cannot do like you have done when you don't use a Keras model. If you don't want to use a model check:
https://www.tensorflow.org/tensorboard/graphs#graphs_of_tffunctions

But as tensorflow/tensorboard#1961 is still open it is better that you still use a model.

sebastian-sz · 2022-05-16T05:22:10Z

@bhack Instead of running the layer, I should wrap it in a tf.keras.Model, compile and run .predict ?

I'm still struggling to create a HLO graph. Will post when I gather more information.

bhack · 2022-05-16T06:54:11Z

@bhack Instead of running the layer, I should wrap it in a tf.keras.Model, compile and run .predict ?

Yes as the graph is on model build:
https://github.com/keras-team/keras/blob/master/keras/engine/training.py#L401-L405

I'm still struggling to create a HLO graph. Will post when I gather more information.

https://www.tensorflow.org/xla#inspect_compiled_programs

LukeWood · 2022-05-16T18:27:29Z

Closes keras-team/keras#422 .

Maybe Apply layer is a wrapper around native Keras layers or BaseImageAugmentationLayer that applies, specified operation to random samples in a batch.

Example usage
images = tf.random.stateless_uniform(shape=(5, 2, 2, 1), seed=[0, 1])

zero_out = tf.keras.layers.Lambda(lambda x: 0 * x)
maybe_apply = MaybeApply(layer=zero_out, rate=0.5, seed=1234)

outputs = maybe_apply(images)
print(outputs)  # Random samples in a batch have been zero'ed out.
Performance / overhead

The layer introduces an overhead for the layer it wraps. It looks like the layer wrapped in MaybeApply takes ~2x the time to execute. I'm not sure if any improvements here can be made - I tried to implement the same behaviour using tf.gather_nd + tf.scatter_nd_update but got similar or even worse results for larger batches.

Below are latency (ms) measurements for Solarization and Posterization. Both were benchmarked with XLA to avoid tf.function retracing.

Known Issues

The layer throws error in XLA with auto_vectorize=True, regardless of what layer it wraps. I'm not really sure why
InvalidArgumentError: Reading input as constant from a dynamic tensor is not yet supported. Xla shape: s32[<=32]
but it works with auto_vectorize=False.

I feel like we need a page in our docs for performance related artifacts like these! These are great to have.

keras_cv/layers/preprocessing/maybe_apply.py

sebastian-sz · 2022-05-17T13:08:42Z

@bhack Following your suggestion on using the tf.keras.Model class I did the benchmarks again using .predict_on_batch and model.compile(jit_compile=...).

When using tf.keras.Model the wrapper layer runs in a similar amount of time as the regular layer.

benchmark_model.zip

bhack · 2022-05-17T13:23:22Z

@sebastian-sz Can you post the same graph without XLA?

As we still don't expose an API to control the XLA compilation x layer in Keras/Keras-cv (see keras-team/keras-io#1541) using just model.compile(jit_compile=...). is very risky cause any single ops not supported by XLA will let the model compilation to fail (see #146 (comment)).

Extra: Without extending Keras-cv layers test for XLA compilation we will never know the exact list of layers with one or more usupported XLA ops.
Also the list/inventory in TF Docs hasn't been updated for years: tensorflow/tensorflow#14798 (comment)

bhack · 2022-05-17T13:27:16Z

P.s. Just to clarify I meant in graph mode as model.compile default but without XLA instread of model.compile eager that I suppose produced the eager graph in your previous post.

Edit:
Checking your ZIP and the XLA boolean I suppose that your 2nd graph is graph mode without XLA instead of eager mode also at it seems to me too fast to be eager. Right?

sebastian-sz · 2022-05-17T14:06:26Z

@bhack

Edit:
Checking your ZIP and the XLA boolean I suppose that your 2nd graph is graph mode without XLA instead of eager mode also at it seems to me too fast to be eager. Right?

Thanks, yes, my bad. Added option for run_eagerly. Eager is slower than graph / XLA (as expected) and the difference between layers is also small:

@sebastian-sz Can you post the same graph without XLA?

I'm not sure I follow - Is the above graph what you requested?

bhack · 2022-05-17T14:24:02Z

I'm not sure I follow - Is the above graph what you requested?

No It was the previous one labeled as eager but It was in graph mode without jit compile

LukeWood · 2022-05-17T17:10:02Z

keras_cv/layers/preprocessing/maybe_apply_test.py

+        with self.assertRaises(ValueError):
+            MaybeApply(rate=invalid_rate, layer=ZeroOut())
+
+    def test_works_with_batched_input(self):


can you pass a seed so that this test is not potentially flaky? Given, it is 1/2^32 flakiness, but still may as well seed it.

Added seed to rng on line 37.

Thanks! does this seed the layer too?

You are right. Added seed param to layer as well.

LukeWood

2 minor comments then good to go.

keras_cv/layers/preprocessing/maybe_apply.py

LukeWood · 2022-05-19T05:45:19Z

Looks good to me @sebastian-sz Thanks for the contribution!

* Added MaybeApply layer. * Changed MaybeApply to override _augment method. * Added seed to maybe_apply_test random generator. * Added seed to layer in batched input test. * Fixed MaybeApply docs.

Silly bug, we were literally just adding the trainable field twice

Added MaybeApply layer.

49b9206

LukeWood reviewed May 16, 2022

View reviewed changes

keras_cv/layers/preprocessing/maybe_apply.py Outdated Show resolved Hide resolved

Changed MaybeApply to override _augment method.

618829b

LukeWood reviewed May 17, 2022

View reviewed changes

LukeWood suggested changes May 17, 2022

View reviewed changes

Added seed to maybe_apply_test random generator.

5136a9d

sebastian-sz requested a review from LukeWood May 17, 2022 17:44

Added seed to layer in batched input test.

39820d4

LukeWood reviewed May 18, 2022

View reviewed changes

keras_cv/layers/preprocessing/maybe_apply.py Show resolved Hide resolved

LukeWood reviewed May 18, 2022

View reviewed changes

keras_cv/layers/preprocessing/maybe_apply.py Outdated Show resolved Hide resolved

LukeWood approved these changes May 18, 2022

View reviewed changes

Fixed MaybeApply docs.

7fa1fe7

LukeWood merged commit 282c66c into keras-team:master May 19, 2022

sebastian-sz deleted the feature-422/add-maybe-apply-layer branch May 19, 2022 05:51

freedomtan pushed a commit to freedomtan/keras-cv that referenced this pull request Jul 20, 2023

Fix extra summary column when show_trainable=True (keras-team#435)

c54f806

Silly bug, we were literally just adding the trainable field twice

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Introduces MaybeApply layer. #435

Introduces MaybeApply layer. #435

sebastian-sz commented May 15, 2022

bhack commented May 15, 2022

sebastian-sz commented May 15, 2022

bhack commented May 15, 2022

sebastian-sz commented May 15, 2022

bhack commented May 15, 2022

bhack commented May 15, 2022 •

edited

Loading

sebastian-sz commented May 16, 2022

bhack commented May 16, 2022

LukeWood commented May 16, 2022

Example usage

Performance / overhead

Known Issues

sebastian-sz commented May 17, 2022

bhack commented May 17, 2022

bhack commented May 17, 2022 •

edited

Loading

sebastian-sz commented May 17, 2022

bhack commented May 17, 2022

LukeWood May 17, 2022

sebastian-sz May 17, 2022

LukeWood May 17, 2022

sebastian-sz May 18, 2022

LukeWood left a comment

LukeWood commented May 19, 2022

Introduces MaybeApply layer. #435

Introduces MaybeApply layer. #435

Conversation

sebastian-sz commented May 15, 2022

Example usage

Performance / overhead

Known Issues

bhack commented May 15, 2022

sebastian-sz commented May 15, 2022

bhack commented May 15, 2022

sebastian-sz commented May 15, 2022

bhack commented May 15, 2022

bhack commented May 15, 2022 • edited Loading

sebastian-sz commented May 16, 2022

bhack commented May 16, 2022

LukeWood commented May 16, 2022

Example usage

Performance / overhead

Known Issues

sebastian-sz commented May 17, 2022

bhack commented May 17, 2022

bhack commented May 17, 2022 • edited Loading

sebastian-sz commented May 17, 2022

bhack commented May 17, 2022

LukeWood May 17, 2022

Choose a reason for hiding this comment

sebastian-sz May 17, 2022

Choose a reason for hiding this comment

LukeWood May 17, 2022

Choose a reason for hiding this comment

sebastian-sz May 18, 2022

Choose a reason for hiding this comment

LukeWood left a comment

Choose a reason for hiding this comment

LukeWood commented May 19, 2022

bhack commented May 15, 2022 •

edited

Loading

bhack commented May 17, 2022 •

edited

Loading