How to add the Attention layer in the TensorFlow GRU Model?

user_15016764 0
Viewed 460

I have created a model for language translation by using TensorFlow Functional API.

This is the model

# encoder 
encoder = tf.keras.Input(shape=(200, ))
enc_embd = tf.keras.layers.Embedding(vocab_train, embedding_dim)(encoder)
encoder_gru = tf.keras.layers.GRU(units, return_sequences=True, return_state=True)
output_e, hidden_e = encoder_gru(enc_embd)

# decoder
decoder = tf.keras.Input(shape=(200, ))
dec_embd = tf.keras.layers.Embedding(vocab_label, embedding_dim)(decoder)
decoder_gru = tf.keras.layers.GRU(units, return_sequences=True, return_state=True)
output_d, hidden_d = decoder_gru(dec_embd, initial_state = hidden_e)
final_output = tf.keras.layers.Dense(vocab_label, activation='softmax')
output_f = final_output(output_d)

I want to ask that, how can I add fully connected tf.keras.layers.Attention (Attention layer) between the encoder and decoder?

keras tensorflow deep-learning machine-learning nlp
1 Answers

You can use the Attention layer between output_e and output_d. Below a full example where we create an autoencoder building a model for encoder and decoder and then merging together.

Define parameters and dummy data:

vocab_train = 111
vocab_label = 123
embedding_dim = 64
units = 32
n_sample = 10
seq_length = 200

X_enc = np.random.randint(0,vocab_train, (n_sample,seq_length))
X_dec = np.random.randint(0,vocab_label, (n_sample,seq_length))
y = np.random.randint(0,2, (n_sample,seq_length,vocab_label))

Define encoder (it must returns also hidden_e because it's used by the decoder):

encoder = tf.keras.Input(shape=(seq_length, ))
enc_embd = tf.keras.layers.Embedding(vocab_train, embedding_dim)(encoder)
encoder_gru = tf.keras.layers.GRU(units, return_sequences=True, return_state=True)
output_e, hidden_e = encoder_gru(enc_embd)

enc = Model(encoder, [hidden_e, output_e])

Define decoder with Attention (it receive as input also output_e and hidden_e):

decoder = tf.keras.Input(shape=(seq_length, ))
hidden_e_input = tf.keras.Input(shape=(units, ))
output_e_input = tf.keras.Input(shape=(seq_length, units))
dec_embd = tf.keras.layers.Embedding(vocab_label, embedding_dim)(decoder)
decoder_gru = tf.keras.layers.GRU(units, return_sequences=True, return_state=True)
output_d, hidden_d = decoder_gru(dec_embd, initial_state = hidden_e_input)
att = tf.keras.layers.Attention()([output_e_input, output_d])
concat = tf.keras.layers.Concatenate()([att, output_d])
final_output = tf.keras.layers.Dense(vocab_label, activation='softmax')(concat)

dec = Model([decoder, hidden_e_input, output_e_input], final_output)

Combine encoder and decoder:

inp_e = tf.keras.Input(shape=(seq_length, ))
h_e, o_e = enc(inp_e)
inp_d = tf.keras.Input(shape=(seq_length, ))
out = dec([inp_d, h_e, o_e])

ae = Model([inp_e, inp_d], out)
ae.compile('adam', 'categorical_crossentropy')
ae.fit([X_enc, X_dec], y, epochs=3)
user_103750490
Related