Stack, shuffle and get back the original tensors-CodePudding

I have following two tensorflow placeholders:

Tensor("Placeholder:0", shape=(32, 2048), dtype=float32)
Tensor("Placeholder:1", shape=(64, 2048), dtype=float32)

Let's call them a and b. I want to stack them and then shuffle randomly. Afterwards, I want to pass them through network. Finally, I want to get the back the a and b before stack and shuffle

What I have done so far

I understand stacking and random shuffle. Therefore, guide me how I can stack them, shuffle them, and finally get back the original indexes.

CodePudding user response：

You can create a shuffle index to the concatenated matrix, so that we know the elements got shuffled went were and then we can put together back using argsort of the indices.

Inputs:

a = tf.random.normal(shape=(32, 2048), dtype=tf.float32)
b = tf.random.normal(shape=(64, 2048), dtype=tf.float32)

Stack the arrays:

c = tf.concat([a,b], axis=0)

Random shuffle:

indices = tf.range(start=0, limit=tf.shape(c)[0], dtype=tf.int32)
shuffled_indices = tf.random.shuffle(indices) #shuffled index will tell where each element of c went.
shuffled_c = tf.gather(c, shuffled_indices)

Get back c, a, b:

getback_c = tf.gather(shuffled_c, tf.argsort(shuffled_indices))
a_1, b_1 = getback_c[:tf.shape(a)[0]], getback_c[tf.shape(a)[0]:]

Verify the values are same:

np.testing.assert_allclose(a.numpy(), a_1.numpy())
np.testing.assert_allclose(b.numpy(), b_1.numpy())