Upload layers.py

layers.py

@@ -0,0 +1,541 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Sep 27 13:11:00 2019
+
+Script for Multi-Head Scaled Dot-Product Attention and Context Attention layers
+
+@author: Chenkai Li
+"""
+
+from tf_keras import backend as K
+from tf_keras import initializers, regularizers, constraints, activations
+from tf_keras.layers import Layer
+
+
+class Attention(Layer):
+    """
+    Adapted from https://github.com/lzfelix/keras_attention
+    
+    Implementation based on the work of Yang et al. "Hierarchical
+    Attention Networks for Document Classification". 
+    [https://www.aclweb.org/anthology/N16-1174]
+    
+    The mathematical formulation of the model is as follows:
+    ```
+    u = f(W * h + b),
+    a_i = softmax(u_i^T * u_s),
+    v_i = \sigma_i a_i * h_i.
+    ```
+    
+    Where h are the input tensors with shape (batch, n_timesteps, hidden_size), for
+    instance, all hidden vectors produced by a recurrent layer, such as a LSTM and the
+    output has shape `(batch, hidden_size)`. This layer also works with inputs with more
+    than 3 dimensions as well, such as sentences in a document, where each input has
+    size (batch, n_docs, n_sentences, embedding_size), outputing 
+    (batch, n_docs, embedding_size)`.
+    
+    
+    # Arguments
+        activation: The activation function f used by the layer (see
+            [activations](../activations.md)). By default tanh is used, another common
+            option is "linear".
+        use_bias: Boolean, whether the layer uses a bias vector.
+        initializer: Initializer for the `kernel` and `context` matrices
+            (see [initializers](../initializers.md)).
+        return_attention: If True, instead of returning the sequence descriptor, this
+            layer will return the computed attention coefficients for each of the
+            sequence timesteps. See Output section for details.
+        W_regularizer: Regularizer function applied to the `kernel` weights matrix
+            (see [regularizer](../regularizers.md)).
+        u_regularizer: Regularizer function applied to the `context` weights matrix
+            (see [regularizer](../regularizers.md)).
+        b_regularizer: Regularizer function applied to the bias vector
+            (see [regularizer](../regularizers.md)).
+        W_constraint: Constraint function applied to the `kernel` weights matrix
+            (see [constraints](../constraints.md)).
+        u_constraint: Constraint function applied to the `contextl` weights matrix
+            (see [constraints](../constraints.md)).
+        b_constraint: Constraint function applied to the bias vector
+            (see [constraints](../constraints.md)).
+
+    # Input shape
+        nD tensor with shape: `(batch_size, ..., timesteps, input_dim)`.
+        The most common situation would be a 3D input with shape
+        `(batch_size, timesteps, input_dim)`.
+
+    # Outuput shape
+        The sequence descriptor with shape `(batch_size, ..., timestamps)`. If
+        `return_attention` is True, this layer will return the `alpha_i` weights
+        for each timestep, and consequently its output shape will be different, namely:
+        `(batch_size, ..., timesteps)`.
+    """
+
+    def __init__(self,
+                 activation='tanh',
+                 initializer='glorot_uniform',
+                 return_attention=False,
+                 W_regularizer=None,
+                 u_regularizer=None,
+                 b_regularizer=None,
+                 W_constraint=None,
+                 u_constraint=None,
+                 b_constraint=None,
+                 bias=True,
+                 **kwargs):
+        
+        self.activation = activations.get(activation)
+        self.initializer = initializers.get(initializer)
+        
+        self.W_regularizer = regularizers.get(W_regularizer)
+        self.u_regularizer = regularizers.get(u_regularizer)
+        self.b_regularizer = regularizers.get(b_regularizer)
+        
+        self.W_constraint = constraints.get(W_constraint)
+        self.u_constraint = constraints.get(u_constraint)
+        self.b_constraint = constraints.get(b_constraint)
+        
+        self.bias = bias
+        self.supports_masking = True
+        self.return_attention = return_attention
+
+        super().__init__(**kwargs)
+
+    def build(self, input_shape):
+        # input_shape: (batch, time, amount_features)
+        
+        # the attention size matches the feature dimension
+        amount_features = input_shape[-1]
+        attention_size  = input_shape[-1]
+
+        self.W = self.add_weight(shape=(amount_features, attention_size),
+                                 initializer=self.initializer,
+                                 regularizer=self.W_regularizer,
+                                 constraint=self.W_constraint,
+                                 name='attention_W')
+        self.b = None
+        if self.bias:
+            self.b = self.add_weight(shape=(attention_size,),
+                                     initializer='zero',
+                                     regularizer=self.b_regularizer,
+                                     constraint=self.b_constraint,
+                                     name='attention_b')
+
+        self.context = self.add_weight(shape=(attention_size,),
+                                       initializer=self.initializer,
+                                       regularizer=self.u_regularizer,
+                                       constraint=self.u_constraint,
+                                       name='attention_us')
+
+        super().build(input_shape)
+
+    def call(self, x, mask=None):        
+        # U = tanh(H*W + b) (eq. 8)        
+        ui = K.dot(x, self.W)              # (b, t, a)
+        if self.b is not None:
+            ui += self.b
+        ui = self.activation(ui)           # (b, t, a)
+
+        # Z = U * us (eq. 9)
+        us = K.expand_dims(self.context)   # (a, 1)
+        ui_us = K.dot(ui, us)              # (b, t, a) * (a, 1) = (b, t, 1)
+        ui_us = K.squeeze(ui_us, axis=-1)  # (b, t, 1) -> (b, t)
+        
+        # alpha = softmax(Z) (eq. 9)
+        alpha = self._masked_softmax(ui_us, mask) # (b, t)
+        alpha = K.expand_dims(alpha, axis=-1)     # (b, t, 1)
+        
+        if self.return_attention:
+            return alpha
+        else:
+            # v = alpha_i * x_i (eq. 10)
+            return K.sum(x * alpha, axis=1)
+    
+    def _masked_softmax(self, logits, mask):
+        """Keras's default implementation of softmax doesn't allow masking, while
+        this method does if `mask` is not `None`."""
+        
+        # softmax(x):
+        #    b = max(x)
+        #    s_i = exp(xi - b) / exp(xj - b)
+        #    return s
+        
+        b = K.max(logits, axis=-1, keepdims=True)
+        logits = logits - b
+
+        exped = K.exp(logits)
+
+        # ignoring masked inputs
+        if mask is not None:
+            mask = K.cast(mask, K.floatx())
+            exped *= mask
+
+        partition = K.sum(exped, axis=-1, keepdims=True)
+
+        # if all timesteps are masked, the partition will be zero. To avoid this
+        # issue we use the following trick:
+        partition = K.maximum(partition, K.epsilon())
+
+        return exped / partition
+
+    def compute_output_shape(self, input_shape):
+        """The attention mechanism computes a weighted average between
+        all hidden vectors generated by the previous sequential layer,
+        hence the input is expected to be
+        `(batch_size, seq_len, amount_features)` if `return_attention` is
+        `False`, otherwise the output should be (batch_size, seq_len)."""
+        if self.return_attention:
+            return input_shape[:-1]
+        else:
+            return input_shape[:-2] + input_shape[-1:]
+
+    def compute_mask(self, x, input_mask=None):
+        """This layer produces a single attended vector from a list
+        of hidden vectors, hence it can't be masked as this means
+        masking a single vector."""
+        return None
+
+    def get_config(self):
+        config = {
+            'activation': self.activation,
+            'initializer': self.initializer,
+            'return_attention': self.return_attention,
+
+            'W_regularizer': initializers.serialize(self.W_regularizer),
+            'u_regularizer': initializers.serialize(self.u_regularizer),
+            'b_regularizer': initializers.serialize(self.b_regularizer),
+
+            'W_constraint': constraints.serialize(self.W_constraint),
+            'u_constraint': constraints.serialize(self.u_constraint),
+            'b_constraint': constraints.serialize(self.b_constraint),
+            
+            'bias': self.bias
+        }
+
+        base_config = super().get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+    
+
+class ScaledDotProductAttention(Layer):
+    """
+    Adapted from https://github.com/CyberZHG/keras-self-attention
+    
+    Implementation based on the work of Vaswani et al. "Attention Is All You 
+    Need". 
+    [https://arxiv.org/pdf/1706.03762.pdf]
+    
+    The attention layer that takes three inputs representing queries, keys and values.
+    \text{Attention}(Q, K, V) = \text{softmax}(\frac{Q K^T}{\sqrt{d_k}}) V
+    """
+
+    def __init__(self,
+                 return_attention=False,
+                 history_only=False,
+                 **kwargs):
+        """Initialize the layer.
+        :param return_attention: Whether to return attention weights.
+        :param history_only: Whether to only use history data.
+        :param kwargs: Arguments for parent class.
+        """
+        super(ScaledDotProductAttention, self).__init__(**kwargs)
+        self.supports_masking = True
+        self.return_attention = return_attention
+        self.history_only = history_only
+
+    def get_config(self):
+        config = {
+            'return_attention': self.return_attention,
+            'history_only': self.history_only,
+        }
+        base_config = super(ScaledDotProductAttention, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+    def compute_output_shape(self, input_shape):
+        if isinstance(input_shape, list):
+            query_shape, key_shape, value_shape = input_shape
+        else:
+            query_shape = key_shape = value_shape = input_shape
+        output_shape = query_shape[:-1] + value_shape[-1:]
+        if self.return_attention:
+            attention_shape = query_shape[:2] + (key_shape[1],)
+            return [output_shape, attention_shape]
+        return output_shape
+
+    def compute_mask(self, inputs, mask=None):
+        if isinstance(mask, list):
+            mask = mask[0]
+        if self.return_attention:
+            return [mask, None]
+        return mask
+
+    def call(self, inputs, mask=None, **kwargs):
+        if isinstance(inputs, list):
+            query, key, value = inputs
+        else:
+            query = key = value = inputs
+        if isinstance(mask, list):
+            mask = mask[1]
+        feature_dim = K.shape(query)[-1]
+        e = K.batch_dot(query, key, axes=2) / K.sqrt(K.cast(feature_dim, dtype=K.floatx()))
+        e = K.exp(e - K.max(e, axis=-1, keepdims=True))
+        if self.history_only:
+            query_len, key_len = K.shape(query)[1], K.shape(key)[1]
+            indices = K.expand_dims(K.arange(0, key_len), axis=0)
+            upper = K.expand_dims(K.arange(0, query_len), axis=-1)
+            e *= K.expand_dims(K.cast(indices <= upper, K.floatx()), axis=0)
+        if mask is not None:
+            e *= K.cast(K.expand_dims(mask, axis=-2), K.floatx())
+        a = e / (K.sum(e, axis=-1, keepdims=True) + K.epsilon())
+        v = K.batch_dot(a, value)
+        if self.return_attention:
+            return [v, a]
+        return v
+
+    
+class MultiHeadAttention(Layer):
+    """
+    Adapted from https://github.com/CyberZHG/keras-multi-head
+    
+    Implementation based on the work of Vaswani et al. "Attention Is All You 
+    Need". 
+    [https://arxiv.org/pdf/1706.03762.pdf]
+    
+    Multi-head scaled dot-product attention layer
+    """
+
+    def __init__(self,
+                 head_num,
+                 activation='relu',
+                 use_bias=True,
+                 kernel_initializer='glorot_normal',
+                 bias_initializer='zeros',
+                 kernel_regularizer=None,
+                 bias_regularizer=None,
+                 kernel_constraint=None,
+                 bias_constraint=None,
+                 history_only=False,
+                 return_multi_attention=False,
+                 **kwargs):
+        """Initialize the layer.
+
+        :param head_num: Number of heads.
+        :param activation: Activations for linear mappings.
+        :param use_bias: Whether to use bias term.
+        :param kernel_initializer: Initializer for linear mappings.
+        :param bias_initializer: Initializer for linear mappings.
+        :param kernel_regularizer: Regularizer for linear mappings.
+        :param bias_regularizer: Regularizer for linear mappings.
+        :param kernel_constraint: Constraints for linear mappings.
+        :param bias_constraint: Constraints for linear mappings.
+        :param history_only: Whether to only use history in attention layer.
+        :param return_multi_attention: Whether to return the attention matrix.
+        """
+        self.supports_masking = True
+        self.head_num = head_num
+        self.activation = activations.get(activation)
+        self.use_bias = use_bias
+        self.kernel_initializer = initializers.get(kernel_initializer)
+        self.bias_initializer = initializers.get(bias_initializer)
+        self.kernel_regularizer = regularizers.get(kernel_regularizer)
+        self.bias_regularizer = regularizers.get(bias_regularizer)
+        self.kernel_constraint = constraints.get(kernel_constraint)
+        self.bias_constraint = constraints.get(bias_constraint)
+        self.history_only = history_only
+        self.return_multi_attention = return_multi_attention
+
+        self.Wq, self.Wk, self.Wv, self.Wo = None, None, None, None
+        self.bq, self.bk, self.bv, self.bo = None, None, None, None
+        super(MultiHeadAttention, self).__init__(**kwargs)
+
+    def get_config(self):
+        config = {
+            'head_num': self.head_num,
+            'activation': activations.serialize(self.activation),
+            'use_bias': self.use_bias,
+            'kernel_initializer': initializers.serialize(self.kernel_initializer),
+            'bias_initializer': initializers.serialize(self.bias_initializer),
+            'kernel_regularizer': regularizers.serialize(self.kernel_regularizer),
+            'bias_regularizer': regularizers.serialize(self.bias_regularizer),
+            'kernel_constraint': constraints.serialize(self.kernel_constraint),
+            'bias_constraint': constraints.serialize(self.bias_constraint),
+            'history_only': self.history_only,
+            'return_multi_attention': self.return_multi_attention
+        }
+        base_config = super(MultiHeadAttention, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+    def compute_output_shape(self, input_shape):
+        if self.return_multi_attention:
+            if isinstance(input_shape, list):
+                q, k, v = input_shape
+                return (q[0], self.head_num, q[1], k[1])
+            return (input_shape[0], self.head_num, input_shape[1], input_shape[1])
+        else:
+            if isinstance(input_shape, list):
+                q, k, v = input_shape
+                return q[:-1] + (v[-1],)
+            return input_shape
+
+    def compute_mask(self, inputs, input_mask=None):
+        if isinstance(input_mask, list):
+            return input_mask[0]
+        return input_mask
+
+    def build(self, input_shape):
+        if isinstance(input_shape, list):
+            q, k, v = input_shape
+        else:
+            q = k = v = input_shape
+        feature_dim = int(v[-1])
+        if feature_dim % self.head_num != 0:
+            raise IndexError('Invalid head number %d with the given input dim %d' % (self.head_num, feature_dim))
+        self.Wq = self.add_weight(
+            shape=(int(q[-1]), feature_dim),
+            initializer=self.kernel_initializer,
+            regularizer=self.kernel_regularizer,
+            constraint=self.kernel_constraint,
+            name='%s_Wq' % self.name,
+        )
+        if self.use_bias:
+            self.bq = self.add_weight(
+                shape=(feature_dim,),
+                initializer=self.bias_initializer,
+                regularizer=self.bias_regularizer,
+                constraint=self.bias_constraint,
+                name='%s_bq' % self.name,
+            )
+        self.Wk = self.add_weight(
+            shape=(int(k[-1]), feature_dim),
+            initializer=self.kernel_initializer,
+            regularizer=self.kernel_regularizer,
+            constraint=self.kernel_constraint,
+            name='%s_Wk' % self.name,
+        )
+        if self.use_bias:
+            self.bk = self.add_weight(
+                shape=(feature_dim,),
+                initializer=self.bias_initializer,
+                regularizer=self.bias_regularizer,
+                constraint=self.bias_constraint,
+                name='%s_bk' % self.name,
+            )
+        self.Wv = self.add_weight(
+            shape=(int(v[-1]), feature_dim),
+            initializer=self.kernel_initializer,
+            regularizer=self.kernel_regularizer,
+            constraint=self.kernel_constraint,
+            name='%s_Wv' % self.name,
+        )
+        if self.use_bias:
+            self.bv = self.add_weight(
+                shape=(feature_dim,),
+                initializer=self.bias_initializer,
+                regularizer=self.bias_regularizer,
+                constraint=self.bias_constraint,
+                name='%s_bv' % self.name,
+            )
+        self.Wo = self.add_weight(
+            shape=(feature_dim, feature_dim),
+            initializer=self.kernel_initializer,
+            regularizer=self.kernel_regularizer,
+            constraint=self.kernel_constraint,
+            name='%s_Wo' % self.name,
+        )
+        if self.use_bias:
+            self.bo = self.add_weight(
+                shape=(feature_dim,),
+                initializer=self.bias_initializer,
+                regularizer=self.bias_regularizer,
+                constraint=self.bias_constraint,
+                name='%s_bo' % self.name,
+            )
+        super(MultiHeadAttention, self).build(input_shape)
+
+    @staticmethod
+    def _reshape_to_batches(x, head_num):
+        input_shape = K.shape(x)
+        batch_size, seq_len, feature_dim = input_shape[0], input_shape[1], input_shape[2]
+        head_dim = feature_dim // head_num
+        x = K.reshape(x, (batch_size, seq_len, head_num, head_dim))
+        x = K.permute_dimensions(x, [0, 2, 1, 3])
+        return K.reshape(x, (batch_size * head_num, seq_len, head_dim))
+
+    @staticmethod
+    def _reshape_from_batches(x, head_num):
+        input_shape = K.shape(x)
+        batch_size, seq_len, feature_dim = input_shape[0], input_shape[1], input_shape[2]
+        x = K.reshape(x, (batch_size // head_num, head_num, seq_len, feature_dim))
+        x = K.permute_dimensions(x, [0, 2, 1, 3])
+        return K.reshape(x, (batch_size // head_num, seq_len, feature_dim * head_num))
+    
+    @staticmethod
+    def _reshape_attention_from_batches(x, head_num):
+        input_shape = K.shape(x)
+        batch_size, seq_len = input_shape[0], input_shape[1]
+        return K.reshape(x, (batch_size // head_num, head_num, seq_len, seq_len))
+
+    @staticmethod
+    def _reshape_mask(mask, head_num):
+        if mask is None:
+            return mask
+        seq_len = K.shape(mask)[1]
+        mask = K.expand_dims(mask, axis=1)
+        mask = K.tile(mask, [1, head_num, 1])
+        return K.reshape(mask, (-1, seq_len))
+
+    def call(self, inputs, mask=None):
+        if isinstance(inputs, list):
+            q, k, v = inputs
+        else:
+            q = k = v = inputs
+        if isinstance(mask, list):
+            q_mask, k_mask, v_mask = mask
+        else:
+            q_mask = k_mask = v_mask = mask
+        q = K.dot(q, self.Wq)
+        k = K.dot(k, self.Wk)
+        v = K.dot(v, self.Wv)
+        if self.use_bias:
+            q += self.bq
+            k += self.bk
+            v += self.bv
+        if self.activation is not None:
+            q = self.activation(q)
+            k = self.activation(k)
+            v = self.activation(v)            
+        y,a = ScaledDotProductAttention(
+            return_attention=True,
+            history_only=self.history_only,
+            name='%s-Attention' % self.name,
+        )(
+            inputs=[
+                self._reshape_to_batches(q, self.head_num),
+                self._reshape_to_batches(k, self.head_num),
+                self._reshape_to_batches(v, self.head_num),
+            ],
+            mask=[
+                self._reshape_mask(q_mask, self.head_num),
+                self._reshape_mask(k_mask, self.head_num),
+                self._reshape_mask(v_mask, self.head_num),
+            ],
+        )
+        
+        y = self._reshape_from_batches(y, self.head_num)
+        a = self._reshape_attention_from_batches(a, self.head_num)
+        y = K.dot(y, self.Wo)
+        if self.use_bias:
+            y += self.bo
+        if self.activation is not None:
+            y = self.activation(y)
+        '''
+        if TF_KERAS:
+            # Add shape information to tensor when using `tf.keras`
+            input_shape = [K.int_shape(q), K.int_shape(k), K.int_shape(v)]
+            output_shape = self.compute_output_shape(input_shape)
+            if output_shape[1] is not None:
+                output_shape = (-1,) + output_shape[1:]
+                y = K.reshape(y, output_shape)
+        '''
+        if self.return_multi_attention:
+            return a
+        return y