Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix ChannelShuffle #418

Closed
wants to merge 4 commits into from
Closed

Fix ChannelShuffle #418

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
69bac3a
alt shuffle indices
innat May 7, 2022
980d9bd
refactor
innat May 7, 2022
06a146c
update
innat May 15, 2022
c103750
Merge branch 'keras-team:master' into fix_chl_sfl
innat May 16, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions keras_cv/layers/preprocessing/channel_shuffle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@

import tensorflow as tf

import keras_cv


@tf.keras.utils.register_keras_serializable(package="keras_cv")
class ChannelShuffle(tf.keras.__internal__.layers.BaseImageAugmentationLayer):
Expand Down Expand Up @@ -52,6 +54,11 @@ def __init__(self, groups=3, seed=None, **kwargs):
self.groups = groups
self.seed = seed

tf.random.set_seed(self.seed)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@LukeWood cc @bhack

This line is only set to satisfy the shuffle-net's requirements. By this, it will produce a static shuffled operation.

aug_fn = ChannelShuffle(groups=3, seed=101)
aug_fn(a, training=True) 

aug_fn = ChannelShuffle(groups=3, seed=42)
aug_fn(a, training=True)

But for random operation, seed=None.
But this may also be a bit in conflict with other KPLs from implementation perspective. Thoughts?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the point in my last comment

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@LukeWood cc @bhack

This line is only set to satisfy the shuffle-net's requirements. By this, it will produce a static shuffled operation.

aug_fn = ChannelShuffle(groups=3, seed=101)
aug_fn(a, training=True) 

aug_fn = ChannelShuffle(groups=3, seed=42)
aug_fn(a, training=True)

But for random operation, seed=None. But this may also be a bit in conflict with other KPLs from implementation perspective. Thoughts?

I don't think we should be modifying seed inside of any layer.

self.rand_uniform = keras_cv.UniformFactorSampler(
lower=0, upper=1, seed=self.seed
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you add unit test for this gradient fix?

)

def augment_image(self, image, transformation=None):
shape = tf.shape(image)
height, width = shape[0], shape[1]
Expand All @@ -65,10 +72,10 @@ def augment_image(self, image, transformation=None):
)

channels_per_group = num_channels // self.groups
image = tf.reshape(image, [height, width, self.groups, channels_per_group])
image = tf.transpose(image, perm=[2, 0, 1, 3])
image = tf.random.shuffle(image, seed=self.seed)
image = tf.transpose(image, perm=[1, 2, 3, 0])

rand_indices = tf.argsort(self.rand_uniform(shape=[self.groups]))
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@LukeWood cc. @bhack
I've noticed that, the tf.argsort fallback to while loop here.

Tensor("loop_body/argsort/TopKV2:1", shape=(3,), dtype=int32)
WARNING:tensorflow:Using a while_loop for converting TopKV2

Copy link
Contributor

@bhack bhack May 16, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@LukeWood cc. @bhack I've noticed that, the tf.argsort fallback to while loop here.

Tensor("loop_body/argsort/TopKV2:1", shape=(3,), dtype=int32)
WARNING:tensorflow:Using a while_loop for converting TopKV2

Just another one to add to the list. But as you can see with the master implementation we have already RandomShuffle: #291 (comment)

So one in, one out.

image = tf.reshape(image, [height, width, channels_per_group, self.groups])
image = tf.gather(image, rand_indices, axis=-1)
image = tf.reshape(image, [height, width, num_channels])

return image
Expand Down