I have a model with multiple layers. I want to shuffle the output tensor of a keras layer before feeding it to the next layer. The case is explained below:

Suppose the output of a layer is:

x = [[[1,2,3], [4,5,6]],
      [[7,8,9], [10,11,12]]]

So, before feeding it into the next layer, I want to randomly shuffle this tensor in a specific way. One example is given below:

x = [[[4,5,6], [1,2,3]]
     [[10,11,12], [7,8,9]]]

I do not want to change the order of individual values. So, the order is changed for x[0] and x[1].

The code of the model is given below. I am also mentioning where I want to shuffle the tensor:

def create_model(embedding_weights, node_vocab_size, path_vocab_size, MAX_SUBTREE_LENGTH):
    config = Config()
    node_input = Input((MAX_SUBTREE_LENGTH,MAX_SUBTREE_LENGTH), dtype=tf.int32)
    path_input = Input((MAX_SUBTREE_LENGTH,), dtype=tf.int32)
    
    #embedding layer
    nodes_embedded = Embedding(node_vocab_size 2, config.embedding_size, trainable = True, name='node_embedding')(node_input)
    path_embedded = Embedding(path_vocab_size 2, config.embedding_size, 
                              trainable = True, name='path_embedding')(path_input) #(b,max_subtree,embedsize)
    
    # path embeddings from node embeddings
    nodes_embedded_merged = K.sum(nodes_embedded, axis=2) #(b,max_subtree,embedsize)
    
    node_path_merged = concatenate([nodes_embedded_merged, path_embedded])
    
    subtree_vectors = TimeDistributed(Dense(config.embedding_size*2, use_bias=False, activation='tanh'))(node_path_merged)
    
    # Here I want to change the subtree_vectors tensor sequence as mentioned above before feeding in to the next layer (the attention layer)
    
    # Attention Layer
    attention_vectors = Dense(1,)(subtree_vectors)
    attention_weights = Softmax(axis=1)(attention_vectors)
    
    # Generating code vectors
    code_vectors = K.sum(subtree_vectors * attention_weights, axis=1)
    
    # Prediction layer
    output_class = Dense(config.num_classes, use_bias=False, activation='softmax')(code_vectors)
    
    model = Model(inputs=[node_input, path_input], outputs=output_class)
    return model

I have tried it in several ways (such as using tf.random.shuffle) but getting errors. I might be having problem with the multiple axis and could not resolve how to do the shuffling in the specific axis. If someone could help me, it would be great! Thanks in advance.

CodePudding user response：

First I am shuffling your inputs according to your need

#First we need to permute according to the axis
#For defining axis
axis = 1 # Use axis 2 if batch-size is included.
perms = list(range(len(subtree_vectors.shape)))
perms[axis], perms[0] = perms[0], perms[axis]

#Define your input x
x = tf.constant([[[1,2,3], [4,5,6]],
      [[7,8,9], [10,11,12]]])

#You need to take the transpose and then shuffle it and then again take the
#transpose then shuffle it then take the transpose, well taking transpose is
#O(1) Operation in Tensorflow.

#In your case set the seed = 6 and axis = 1

tf.random.set_seed(6)
shuffled_x = tf.transpose(tf.random.shuffle(tf.transpose(tf.random.shuffle(x), perm=perms)), perm=perms)

<tf.Tensor: shape=(2, 2, 3), dtype=int32, numpy=
array([[[ 4,  5,  6],
        [ 1,  2,  3]],

       [[10, 11, 12],
        [ 7,  8,  9]]])>

Now, you need to apply this operation to your output

subtree_vectors = TimeDistributed(Dense(3, use_bias=False, activation='tanh'))(x)

tf.random.set_seed(6)
tf.keras.layers.Lambda(lambda x : tf.transpose(tf.random.shuffle(tf.transpose(tf.random.shuffle(x), perm=perms)), perm=perms))(subtree_vectors)

Output Before Shuffling

<tf.Tensor: shape=(2, 2, 3), dtype=float32, numpy=
array([[[-0.95375925, -0.9924454 ,  0.63576657],
        [-0.94918346, -0.99999785,  0.9297005 ]],

       [[-0.94416803, -1.        ,  0.9881501 ],
        [-0.93867296, -1.        ,  0.99805164]]], dtype=float32)>

Output After Shuffling

<tf.Tensor: shape=(2, 2, 3), dtype=float32, numpy=
array([[[-0.94918346, -0.99999785,  0.9297005 ],
        [-0.95375925, -0.9924454 ,  0.63576657]],

       [[-0.93867296, -1.        ,  0.99805164],
        [-0.94416803, -1.        ,  0.9881501 ]]], dtype=float32)>