Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- 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[2]))
- # output shape == (batch_size, vocab)
- x = self.fc(output)
- return x, state, attention_weights
- decoder = Decoder(vocab_out_size, embedding_dim, units, BATCH_SIZE)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement