Luong Style Attention Mechanism with Dot and General scoring functions in keras and tensorflow

Viewed 752

I am trying to implement the dot product and general implementation of calculating similarity scores from encoder and decoder output and hidden states respectively in keras.

I have got the idea to do the product of tf.keras.layers.dot(encoder_output,decoder_state) for calculating product score but there is error in multiplication of these two values.

class Attention(tf.keras.Model):
  def __init__(self,units):
    super().__init__()
    self.units = units

  def call(self, decoder_state, encoder_output):
      score = tf.keras.layers.dot([encoder_output,decoder_state], axes=[2, 1])
      
      attention_weights = tf.nn.softmax(score, axis=1)
      context_vector = attention_weights * encoder_output
      context_vector = tf.reduce_sum(context_vector, axis=1)
      
      return context_vector, attention_weights

batch_size = 16
units = 32
input_length = 20
decoder_state = tf.random.uniform(shape=[batch_size, units])
encoder_output = tf.random.uniform(shape=[batch_size, input_length, units])
attention = Attention(units)
context_vector, attention_weights = attention(decoder_state, encoder_output)

I am getting the following error:

/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: Incompatible shapes: [16,20] vs. [16,20,32] [Op:Mul]

It is a very simple fix but as I am new to this I am not able to get the exact method needed to be called here. I have tried reshaping the values of encoder_output but still this does not work. Request to help me fix this.

1 Answers

I am just putting @Ayush Srivastava's comment as a response so that the post gets an answer.

Basically, the error occurs because you are trying to multiply 2 tensors (namely attention_weights and encoder_output) with different shapes, so you need to reshape the decoder_state.

Here is the full answer:

class Attention(tf.keras.Model):
    def __init__(self,units):
        super().__init__()
        self.units = units

    def call(self, decoder_state, encoder_output):
        decoder_state = tf.keras.layers.Reshape((decoder_state.shape[1], 1))(decoder_state) 
        score = tf.keras.layers.dot([encoder_output, decoder_state],[2, 1]) 

        attention_weights = tf.nn.softmax(score, axis=1)

        context_vector = attention_weights * encoder_output
        context_vector = tf.reduce_sum(context_vector, axis=1)
        
        return context_vector, attention_weights

Shapes:

decoder_state before reshape: (16, 32)
decoder_state after reshape:  (16, 32, 1)
enc_output:                   (16, 20, 32)
score:                        (16, 20, 1)
attention_weights:            (16, 20, 1)
context_vector before sum:    (16, 20, 32)
Related