Vectorize ChannelShuffle #1433

james77777778 · 2023-02-22T03:03:43Z

What does this PR do?

Benchmark results:

Runtime gets high at some points due to retracing and some warning.

WARNING:tensorflow:Using a while_loop for converting RandomShuffle cause there is no registered converter for this op.

Here is the log.txt

Overall, I think vectorized implementation beats old one.

To check the consistency, I manually set batch_size=1 to make tf.random.shuffle acts the same for new and old implementation:

class ChannelShuffleTest(tf.test.TestCase):
    def test_consistency_with_old_impl(self):
        # must set batch_size=1 due to randomness from
        # images = tf.random.shuffle(images, seed=self.seed)
        image_shape = (1, 32, 32, 3)
        fixed_seed = 2023
        image = tf.random.uniform(shape=image_shape) * 255.0

        layer = ChannelShuffle(groups=3, seed=fixed_seed)
        old_layer = OldChannelShuffle(groups=3, seed=fixed_seed)

        output = layer(image)
        old_output = old_layer(image)

        self.assertAllClose(old_output, output)

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

james77777778 · 2023-02-22T07:44:19Z

In this PR, I find that channel shuffle cannot process independently.
Need some time to survey but it could be difficult because tf.random.shuffle cannot apply on batched inputs.

LukeWood · 2023-02-22T08:40:09Z

In this PR, I find that channel shuffle cannot process independently. Need some time to survey but it could be difficult because tf.random.shuffle cannot apply on batched inputs.

can we roll our own version instead of using tf.random.shuffle? I believe this could be done with tf.gather_nd() and tf.random.uniform().

james77777778 · 2023-02-22T08:48:54Z

In this PR, I find that channel shuffle cannot process independently. Need some time to survey but it could be difficult because tf.random.shuffle cannot apply on batched inputs.

can we roll our own version instead of using tf.random.shuffle? I believe this could be done with tf.gather_nd() and tf.random.uniform().

Thanks for the keywords.
I have tested tf.TensorArray (as a output list) with similar implementation from previous one but the performance decreased a lot...

I'm going to try tf.gather_nd() and tf.random.uniform()

bhack · 2023-02-22T10:41:27Z

I think we can do:

num_channels = x.shape.as_list()[-1]

# Generate a random permutation of the channel indices
shuffled_indices = tf.random.shuffle(tf.range(num_channels), seed=seed)

# Permute the channels of the input tensor according to the random indices
shuffled_x = tf.gather(x, shuffled_indices, axis=-1)

P.s. We could also use the stateless version:
tf.random.experimental.stateless_shuffle

- replace tf.random.shuffle by tf.random.uniform and tf.argsort - add test of independence - remove numerical check

james77777778 · 2023-02-23T02:56:18Z

@LukeWood @bhack
Thanks for your guidance

Much of the difficulty comes from how to independently shuffle indices in batch.

I adopt the rand+argsort trick based on
https://stackoverflow.com/a/55317373

    def get_random_transformation_batch(self, batch_size, **kwargs):
        # get batched shuffled indices
        # for example: batch_size=2; self.group=5
        # indices = [
        #     [0, 2, 3, 4, 1],
        #     [4, 1, 0, 2, 3]
        # ]
        indices_distribution = tf.random.uniform(
            (batch_size, self.groups), seed=self.seed
        )
        indices = tf.argsort(indices_distribution, axis=-1)
        return indices

and then gather by

    def augment_images(self, images, transformations, **kwargs):
        ......
        indices = transformations

        # append batch indexes next to shuffled indices
        batch_indexs = tf.reshape(
            tf.repeat(tf.range(batch_size), self.groups),
            (batch_size, self.groups),
        )
        indices = tf.stack([batch_indexs, indices], axis=-1)

        channels_per_group = num_channels // self.groups
        images = tf.reshape(
            images, [batch_size, height, width, self.groups, channels_per_group]
        )
        images = tf.transpose(images, perm=[0, 3, 1, 2, 4])
        images = tf.gather_nd(images, indices=indices)
        images = tf.transpose(images, perm=[0, 2, 3, 4, 1])
        images = tf.reshape(images, [batch_size, height, width, num_channels])
        return images

All unit tests passed but it is difficult to check the consistency with non-vectorized implementation because the randomness comes from different sources (tf.random.shuffle, tf.random.uniform).
I just removed the test in benchmark script.

Also a independence test is added.

The performance result:

bhack · 2023-02-23T03:01:05Z

Much of the difficulty comes from how to independently shuffle indices in batch.

If I remember correctly is tf.random.experimental.stateless_shuffle

james77777778 · 2023-02-23T03:27:11Z

Much of the difficulty comes from how to independently shuffle indices in batch.

If I remember correctly is tf.random.experimental.stateless_shuffle

I think tf.random.experimental.stateless_shuffle acts the same as tf.random.shuffle
So they can only shuffle tensor along 1 axis?

In ChannelShuffle, we want to shuffle channels independently in each image like:
idx=1 => shuffled_idx = [1, 2, 0]
idx=2 => shuffled_idx = [0, 2, 1]
...

It is achievable to loop over each image with the operations you mentioned.
Pseudo code like:

outputs = tf.TensorArray(...)
for i in range(batch_size):
    num_channels = images[i].shape.as_list()[-1]
    shuffled_indices = tf.random.shuffle(tf.range(num_channels), seed=seed)
    shuffled_x = tf.gather(images[i], shuffled_indices, axis=-1)
    outputs.write(i, shuffled_x)
outputs.stack()

But it is too slow that it is even slower than non-vectorized version in eager mode.

I cannot figure out how to vectorize this with tf.random.shuffle or tf.random.experimental.stateless_shuffle. Please enlighten me if I missed something.

bhack · 2023-02-23T03:50:26Z

In ChannelShuffle, we want to shuffle channels independently in each image like:
idx=1 => shuffled_idx = [1, 2, 0]
idx=2 => shuffled_idx = [0, 2, 1]

Something like this?
https://colab.research.google.com/gist/bhack/65760a7896ee9f6ec450ebbe8bd32ca2/exp_random_shuffle.ipynb

james77777778 · 2023-02-23T03:59:49Z

In ChannelShuffle, we want to shuffle channels independently in each image like:
idx=1 => shuffled_idx = [1, 2, 0]
idx=2 => shuffled_idx = [0, 2, 1]

Something like this? https://colab.research.google.com/gist/bhack/65760a7896ee9f6ec450ebbe8bd32ca2/exp_random_shuffle.ipynb

https://colab.research.google.com/drive/1b-tZErA-jCv6SNurdprdOYc6dLFr0ZSH?usp=sharing
I manually set the values to show that

only 1 shuffled indices within the batch
outputs are shuffled by same shuffled indices and not independently across the batch

bhack · 2023-02-23T04:05:24Z

Ok, so you meant within the batch. Sorry I was confused what kind of randomization policy we finally selected for the new vectorizing refactoring (mandatory within the batch randomization or not)

james77777778 · 2023-03-07T02:45:29Z

Kindly ping @LukeWood @ianstenbit

The idea of vectorizing please see #1433 (comment)

The current implementation shuffles channel (by groups) independently for each image in the batch

ianstenbit · 2023-03-07T16:45:48Z

@LukeWood PTAL -- what's the decision re: independent augmentation in batches?

LukeWood · 2023-03-07T20:16:08Z

All augmentation should be done independently on a per-image basis.

LukeWood · 2023-03-07T20:16:47Z

/gcbrun

Thanks you @james77777778 !!!

LukeWood · 2023-03-14T18:37:27Z

Sorry one more thing: Can add a test like the one in https://github.com/keras-team/keras-cv/pull/1480/files#diff-1c771135874332a25b437da4a5706faea85ff719331dfe7beb255403a1a66477

james77777778 · 2023-03-15T00:57:25Z

Sorry one more thing: Can add a test like the one in https://github.com/keras-team/keras-cv/pull/1480/files#diff-1c771135874332a25b437da4a5706faea85ff719331dfe7beb255403a1a66477

Hi @LukeWood
It should be already added at
https://github.com/keras-team/keras-cv/pull/1433/files#diff-67ea20e78d28e97daedb7b45698d50914d1ffef26cc62311c855606d6d171fc0R98-R105

EDITED:
I have added other tests from RandomTranslation including
test_config_with_custom_name, test_output_dtypes, test_config

https://github.com/keras-team/keras-cv/pull/1433/files#diff-67ea20e78d28e97daedb7b45698d50914d1ffef26cc62311c855606d6d171fc0R98-R123

- `test_config_with_custom_name` - `test_output_dtypes` - `test_config`

ianstenbit

Thanks for the PR!

ianstenbit · 2023-03-16T14:46:48Z

benchmarks/vectorized_channel_shuffle.py

+)
+
+
+class OldChannelShuffle(BaseImageAugmentationLayer):


Unless the behavior is changing as part of this PR, we should add a PyTest to this benchmark to verify that the old/new implementations are numerically identical (like this one)

@LukeWood to confirm

Please see #1433 (comment)

Great -- thank you!

james77777778 · 2023-03-17T01:48:34Z

Hi @ianstenbit

Luckily, I got a magic number (seed=2023 😂) that has the same randomness for old/new implementation. I have added the unit test to the benchmark.

Also I ran the script examples/layers/preprocessing/classification/channel_shuffle_demo.py to visualize the output of new ChannelShuffle

ianstenbit · 2023-03-17T17:51:44Z

/gcbrun

ianstenbit · 2023-03-17T19:10:03Z

GCB CI tests failing due to unrelated changed, should be green post-merge.

* Vectorize ChannelShuffle * Update ChannelShuffle. - replace tf.random.shuffle by tf.random.uniform and tf.argsort - add test of independence - remove numerical check * Fix typo * add tests. - `test_config_with_custom_name` - `test_output_dtypes` - `test_config` * Add unit test to verify numerical consistency * Fix ChannelShuffle

Vectorize ChannelShuffle

f7e2c14

Update ChannelShuffle.

a7c1c30

- replace tf.random.shuffle by tf.random.uniform and tf.argsort - add test of independence - remove numerical check

Fix typo

f8f0710

ianstenbit requested a review from LukeWood March 3, 2023 19:40

LukeWood approved these changes Mar 7, 2023

View reviewed changes

james77777778 mentioned this pull request Mar 10, 2023

vectorize KPLS: tf.vectorized_map fallback cause #291

Closed

james77777778 added 2 commits March 15, 2023 09:11

Merge branch 'keras-team:master' into channel-shuffle

4c21217

add tests.

56c7a83

- `test_config_with_custom_name` - `test_output_dtypes` - `test_config`

james77777778 requested a review from LukeWood March 16, 2023 03:04

ianstenbit approved these changes Mar 16, 2023

View reviewed changes

james77777778 added 2 commits March 17, 2023 01:41

Add unit test to verify numerical consistency

503e6c4

Fix ChannelShuffle

d3fcaa2

ianstenbit merged commit 10aac22 into keras-team:master Mar 17, 2023

james77777778 deleted the channel-shuffle branch March 18, 2023 05:37

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Vectorize ChannelShuffle #1433

Vectorize ChannelShuffle #1433

james77777778 commented Feb 22, 2023

james77777778 commented Feb 22, 2023

LukeWood commented Feb 22, 2023

james77777778 commented Feb 22, 2023

bhack commented Feb 22, 2023 •

edited

Loading

james77777778 commented Feb 23, 2023

bhack commented Feb 23, 2023

james77777778 commented Feb 23, 2023

bhack commented Feb 23, 2023 •

edited

Loading

james77777778 commented Feb 23, 2023 •

edited

Loading

bhack commented Feb 23, 2023 •

edited

Loading

james77777778 commented Mar 7, 2023

ianstenbit commented Mar 7, 2023

LukeWood commented Mar 7, 2023

LukeWood commented Mar 7, 2023

LukeWood commented Mar 14, 2023

james77777778 commented Mar 15, 2023 •

edited

Loading

ianstenbit left a comment

ianstenbit Mar 16, 2023

ianstenbit Mar 16, 2023

james77777778 Mar 17, 2023

ianstenbit Mar 17, 2023

james77777778 commented Mar 17, 2023

ianstenbit commented Mar 17, 2023

ianstenbit commented Mar 17, 2023

Vectorize ChannelShuffle #1433

Vectorize ChannelShuffle #1433

Conversation

james77777778 commented Feb 22, 2023

What does this PR do?

Before submitting

Who can review?

james77777778 commented Feb 22, 2023

LukeWood commented Feb 22, 2023

james77777778 commented Feb 22, 2023

bhack commented Feb 22, 2023 • edited Loading

james77777778 commented Feb 23, 2023

bhack commented Feb 23, 2023

james77777778 commented Feb 23, 2023

bhack commented Feb 23, 2023 • edited Loading

james77777778 commented Feb 23, 2023 • edited Loading

bhack commented Feb 23, 2023 • edited Loading

james77777778 commented Mar 7, 2023

ianstenbit commented Mar 7, 2023

LukeWood commented Mar 7, 2023

LukeWood commented Mar 7, 2023

LukeWood commented Mar 14, 2023

james77777778 commented Mar 15, 2023 • edited Loading

ianstenbit left a comment

Choose a reason for hiding this comment

ianstenbit Mar 16, 2023

Choose a reason for hiding this comment

ianstenbit Mar 16, 2023

Choose a reason for hiding this comment

james77777778 Mar 17, 2023

Choose a reason for hiding this comment

ianstenbit Mar 17, 2023

Choose a reason for hiding this comment

james77777778 commented Mar 17, 2023

ianstenbit commented Mar 17, 2023

ianstenbit commented Mar 17, 2023

bhack commented Feb 22, 2023 •

edited

Loading

bhack commented Feb 23, 2023 •

edited

Loading

james77777778 commented Feb 23, 2023 •

edited

Loading

bhack commented Feb 23, 2023 •

edited

Loading

james77777778 commented Mar 15, 2023 •

edited

Loading