tensorflow - Custom Attention Layer using in Keras

Question

I want to create a custom attention layer that for input at any time this layer returns the weighted mean of inputs at all time inputs.

For Example, I want that input tensor with shape [32,100,2048] goes to layer and I get the tensor with the shape [32,100,2048]. I wrote the Layer as follow:

import tensorflow as tf

from keras.layers import Layer, Dense

#or

from tensorflow.keras.layers import Layer, Dense


class Attention(Layer):

  def __init__(self, units_att):

     self.units_att = units_att
     self.W = Dense(units_att)
     self.V = Dense(1)
     super().__init__()

  def __call__(self, values):

      t = tf.constant(0, dtype= tf.int32)    
      time_steps = tf.shape(values)[1]
      initial_outputs = tf.TensorArray(dtype=tf.float32, size=time_steps)
      initial_att =  tf.TensorArray(dtype=tf.float32, size=time_steps)

      def should_continue(t, *args):
          return t < time_steps

      def iteration(t, values, outputs, atts):

        score = self.V(tf.nn.tanh(self.W(values)))

        # attention_weights shape == (batch_size, time_step, 1)
        attention_weights = tf.nn.softmax(score, axis=1)

        # context_vector shape after sum == (batch_size, hidden_size)
        context_vector = attention_weights * values
        context_vector = tf.reduce_sum(context_vector, axis=1)

        outputs = outputs.write(t, context_vector)
        atts = atts.write(t, attention_weights)
        return t + 1, values, outputs, atts

      t, values, outputs, atts = tf.while_loop(should_continue, iteration,
                                  [t, values, initial_outputs, initial_att])

      outputs = outputs.stack()
      outputs = tf.transpose(outputs, [1,0,2])

      atts = atts.stack()
      atts = tf.squeeze(atts, -1)
      atts = tf.transpose(atts, [1,0,2])
      return t, values, outputs, atts

For input= tf.constant(2, shape= [32, 100, 2048], dtype= tf.float32) I get the output with shape = [32,100,2048] in tf2 and [32,None, 2048] in tf1.

For Input input= Input(shape= (None, 2048)) I get the output with shape = [None, None, 2048] in tf1 and I get error

TypeError: 'Tensor' object cannot be interpreted as an integer

in tf2.

Finally, in both cases, I can't use this layer in my model because my model input is Input(shape= (None, 2048)) and I get the error

AttributeError: 'NoneType' object has no attribute '_inbound_nodes'

in tf1 and in tf2 I get the same error as said in above, I create my model with Keras functional method.

Answer 1

From the code you have shared, looks like you want to implement Bahdanau's attention layer in your code. You want to attend to all the 'values' (prev layer output - all its hidden states) and your 'query' would be the last hidden state of the decoder. Your code should actually be very simple and should look like:

        class Bahdanau(tf.keras.layers.Layer):
            def __init__(self, n):
                super(Bahdanau, self).__init__()
                self.w = tf.keras.layers.Dense(n)
                self.u = tf.keras.layers.Dense(n)
                self.v = tf.keras.layers.Dense(1)
        
            def call(self, query, values):
                query = tf.expand_dims(query, 1)
                e = self.v(tf.nn.tanh(self.w(query) + self.u(values)))
                a = tf.nn.softmax(e, axis=1)
                c = a * values
                c = tf.reduce_sum(c, axis=1)
                return a,c
        
        ##Say we want 10 units in the single layer MLP determining w,u
        attentionlayer = Bahdanau(10)
        ##Call with i/p: decoderstate @ t-1 and all encoder hidden states
        a, c = attentionlayer(stminus1, hj)
    

We are not specifying the tensor shape anywhere in the code. This code will return you a context tensor of same size as 'stminus1' which is the 'query'. It does this after attending to all the 'values' (all output states of decoder) using Bahdanau's attention mechanism.

So assuming your batch size is 32, timesteps=100 and embedding dimension=2048, the shape of stminus1 should be (32,2048) and the shape of the hj should be (32,100,2048). The shape of the output context would be (32,2048). We also returned the 100 attention weights just in case you want to route them to a nice display.

This is the simplest version of 'Attention'. If you have any other intent, please let me know and I will reformat my answer. For more specific details, please refer https://towardsdatascience.com/create-your-own-custom-attention-layer-understand-all-flavours-2201b5e8be9e