On the official tensorflow page there is one exmple of a decoder (https://www.tensorflow.org/tutorials/text/nmt_with_attention#next_steps):
class Decoder(tf.keras.Model): def __init__(self, vocab_size, embedding_dim, dec_units, batch_sz): super(Decoder, self).__init__() self.batch_sz = batch_sz self.dec_units = dec_units self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim) self.gru = tf.keras.layers.GRU(self.dec_units, return_sequences=True, return_state=True, recurrent_initializer='glorot_uniform') self.fc = tf.keras.layers.Dense(vocab_size) # used for attention self.attention = BahdanauAttention(self.dec_units) def call(self, x, hidden, enc_output): # enc_output shape == (batch_size, max_length, hidden_size) context_vector, attention_weights = self.attention(hidden, enc_output) # x shape after passing through embedding == (batch_size, 1, embedding_dim) x = self.embedding(x) # x shape after concatenation == (batch_size, 1, embedding_dim + hidden_size) x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1) # passing the concatenated vector to the GRU output, state = self.gru(x) # output shape == (batch_size * 1, hidden_size) output = tf.reshape(output, (-1, output.shape)) # output shape == (batch_size, vocab) x = self.fc(output) return x, state, attention_weights
The output here is of shape (batch_size, vocab). I would like to have an output with shape (batch_size, max_length, vocab), where max_length is the output sequence length. I suppose there can be something done with the TimeDistributed layer but I tried several things and nothing realy worked out. Is there any work around to obtain this?
One way would be:
self.attention = layers.TimeDristributed(BahdanauAttention(self.dec_units))
but since BahdanauAttention has two inputs I do not know how to fix it, since timedistributed layer does not work with multiple inputs easily.