Train_MTL_with_report_gen.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.Encoder import Encoder\n",
    "from utils.Classifier import Classifier\n",
    "from utils.Detector import Detector\n",
    "from utils.Segmenter import Segmenter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import keras.backend as K\n",
    "from keras.layers import Layer, Dense, TimeDistributed, Concatenate, InputSpec,  RNN\n",
    "from keras.layers.wrappers import Wrapper\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "class ScaledDotProductAttention(Layer):\n",
    "    \"\"\"\n",
    "        Implementation according to:\n",
    "            \"Attention is all you need\" by A Vaswani, N Shazeer, N Parmar (2017)\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, return_attention=False, **kwargs):    \n",
    "        self._return_attention = return_attention\n",
    "        self.supports_masking = True\n",
    "        super(ScaledDotProductAttention, self).__init__(**kwargs)\n",
    "    \n",
    "    def compute_output_shape(self, input_shape):\n",
    "        self._validate_input_shape(input_shape)\n",
    "\n",
    "        if not self._return_attention:\n",
    "            return input_shape[-1]\n",
    "        else:\n",
    "            return [input_shape[-1], [input_shape[0][0], input_shape[0][1], input_shape[1][2]]]\n",
    "    \n",
    "    def _validate_input_shape(self, input_shape):\n",
    "        if len(input_shape) != 3:\n",
    "            raise ValueError(\"Layer received an input shape {0} but expected three inputs (Q, V, K).\".format(input_shape))\n",
    "        else:\n",
    "            if input_shape[0][0] != input_shape[1][0] or input_shape[1][0] != input_shape[2][0]:\n",
    "                raise ValueError(\"All three inputs (Q, V, K) have to have the same batch size; received batch sizes: {0}, {1}, {2}\".format(input_shape[0][0], input_shape[1][0], input_shape[2][0]))\n",
    "            if input_shape[0][1] != input_shape[1][1] or input_shape[1][1] != input_shape[2][1]:\n",
    "                raise ValueError(\"All three inputs (Q, V, K) have to have the same length; received lengths: {0}, {1}, {2}\".format(input_shape[0][0], input_shape[1][0], input_shape[2][0]))\n",
    "            if input_shape[0][2] != input_shape[1][2]:\n",
    "                raise ValueError(\"Input shapes of Q {0} and V {1} do not match.\".format(input_shape[0], input_shape[1]))\n",
    "    \n",
    "    def build(self, input_shape):\n",
    "        self._validate_input_shape(input_shape)\n",
    "        \n",
    "        super(ScaledDotProductAttention, self).build(input_shape)\n",
    "    \n",
    "    def call(self, x, mask=None):\n",
    "        q, k, v = x\n",
    "        d_k = q.shape.as_list()[2]\n",
    "\n",
    "        # in pure tensorflow:\n",
    "        # weights = tf.matmul(x_batch, tf.transpose(y_batch, perm=[0, 2, 1]))\n",
    "        # normalized_weights = tf.nn.softmax(weights/scaling)\n",
    "        # output = tf.matmul(normalized_weights, x_batch)\n",
    "        \n",
    "        weights = K.batch_dot(q,  k, axes=[2, 2])\n",
    "\n",
    "        if mask is not None:\n",
    "            # add mask weights\n",
    "            if isinstance(mask, (list, tuple)):\n",
    "                if len(mask) > 0:\n",
    "                    raise ValueError(\"mask can only be a Tensor or a list of length 1 containing a tensor.\")\n",
    "\n",
    "                mask = mask[0]\n",
    "\n",
    "            weights += -1e10*(1-mask)\n",
    "\n",
    "        normalized_weights = K.softmax(weights / np.sqrt(d_k))\n",
    "        output = K.batch_dot(normalized_weights, v)\n",
    "        \n",
    "        if self._return_attention:\n",
    "            return [output, normalized_weights]\n",
    "        else:\n",
    "            return output\n",
    "\n",
    "    def get_config(self):\n",
    "        config = {'return_attention': self._return_attention}\n",
    "        base_config = super(ScaledDotProductAttention, self).get_config()\n",
    "        return dict(list(base_config.items()) + list(config.items()))\n",
    "\n",
    "class MultiHeadAttention():\n",
    "    \"\"\"\n",
    "        Implementation according to:\n",
    "            \"Attention is all you need\" by A Vaswani, N Shazeer, N Parmar (2017)\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, h, d_k=None, d_v=None, d_model=None, activation=None, return_attention=False, **kwargs):    \n",
    "        super(MultiHeadAttention, self).__init__(**kwargs)\n",
    "        \n",
    "        if (type(h) is not int or h < 2):\n",
    "            raise ValueError(\"You have to set `h` to an int >= 2.\")\n",
    "        self._h = h\n",
    "        \n",
    "        if d_model and (type(d_model) is not int or d_model < 1):\n",
    "                raise ValueError(\"You have to set `d_model` to an int >= 1.\")\n",
    "        self._d_model = d_model\n",
    "        \n",
    "        if d_k and int (type(d_k) is not int or d_k < 1):\n",
    "            raise ValueError(\"You have to set `d_k` to an int >= 1.\")\n",
    "        self._d_k = d_k\n",
    "        \n",
    "        if d_v and (type(d_v) is not int or d_v < 1):\n",
    "            raise ValueError(\"You have to set `d_v` to an int >= 1.\")\n",
    "        self._d_v = d_v\n",
    "        \n",
    "        self._activation = None\n",
    "        self._return_attention = return_attention\n",
    "    \n",
    "    def compute_output_shape(self, input_shape):\n",
    "        self._validate_input_shape(input_shape)\n",
    "        \n",
    "        if self._return_attention:\n",
    "            return [input_shape[-1], [input_shape[0][0], input_shape[1][1], self._h*input_shape[2][2]]]\n",
    "        else:\n",
    "            return input_shape[-1]\n",
    "    \n",
    "    def _validate_input_shape(self, input_shape):\n",
    "        if len(input_shape) != 3:\n",
    "            raise ValueError(\"Layer received an input shape {0} but expected three inputs (Q, V, K).\".format(input_shape))\n",
    "        else:\n",
    "            if input_shape[0][0] != input_shape[1][0] or input_shape[1][0] != input_shape[2][0]:\n",
    "                raise ValueError(\"All three inputs (Q, V, K) have to have the same batch size; received batch sizes: {0}, {1}, {2}\".format(input_shape[0][0], input_shape[1][0], input_shape[2][0]))\n",
    "            if input_shape[0][1] != input_shape[1][1] or input_shape[1][1] != input_shape[2][1]:\n",
    "                raise ValueError(\"All three inputs (Q, V, K) have to have the same length; received lengths: {0}, {1}, {2}\".format(input_shape[0][0], input_shape[1][0], input_shape[2][0]))\n",
    "            if input_shape[0][2] != input_shape[1][2]:\n",
    "                raise ValueError(\"Input shapes of Q {0} and V {1} do not match.\".format(input_shape[0], input_shape[1]))\n",
    "    \n",
    "    def build(self, input_shape):\n",
    "        self._validate_input_shape(input_shape)\n",
    "        \n",
    "        d_k = self._d_k if self._d_k else input_shape[1][-1]\n",
    "        d_model = self._d_model if self._d_model else input_shape[1][-1]\n",
    "        d_v = self._d_v\n",
    "\n",
    "        if type(d_k) == tf.Dimension:\n",
    "            d_k = d_k.value\n",
    "        if type(d_model) == tf.Dimension:\n",
    "            d_model = d_model.value\n",
    "        \n",
    "        self._q_layers = []\n",
    "        self._k_layers = []\n",
    "        self._v_layers = []\n",
    "        self._sdp_layer = ScaledDotProductAttention(return_attention=self._return_attention)\n",
    "    \n",
    "        for _ in range(self._h):\n",
    "            self._q_layers.append(\n",
    "                TimeDistributed(\n",
    "                    Dense(d_k, activation=self._activation, use_bias=False)\n",
    "                )\n",
    "            )\n",
    "            self._k_layers.append(\n",
    "                TimeDistributed(\n",
    "                    Dense(d_k, activation=self._activation, use_bias=False)\n",
    "                )\n",
    "            )\n",
    "            self._v_layers.append(\n",
    "                TimeDistributed(\n",
    "                    Dense(d_v, activation=self._activation, use_bias=False)\n",
    "                )\n",
    "            )\n",
    "        \n",
    "        self._output = TimeDistributed(Dense(d_model))\n",
    "        #if self._return_attention:\n",
    "        #    self._output = Concatenate()\n",
    "    \n",
    "    def __call__(self, x, mask=None):\n",
    "        if isinstance(x, (list, tuple)):\n",
    "            self.build([it.shape for it in x])\n",
    "        else:\n",
    "            self.build(x.shape)\n",
    "\n",
    "        q, k, v = x\n",
    "        \n",
    "        outputs = []\n",
    "        attentions = []\n",
    "        for i in range(self._h):\n",
    "            qi = self._q_layers[i](q)\n",
    "            ki = self._k_layers[i](k)\n",
    "            vi = self._v_layers[i](v)\n",
    "            \n",
    "            if self._return_attention:\n",
    "                output, attention = self._sdp_layer([qi, ki, vi], mask=mask)\n",
    "                outputs.append(output)\n",
    "                attentions.append(attention)\n",
    "            else:\n",
    "                output = self._sdp_layer([qi, ki, vi], mask=mask)\n",
    "                outputs.append(output)\n",
    "            \n",
    "        concatenated_outputs = Concatenate()(outputs)\n",
    "        output = self._output(concatenated_outputs)\n",
    "        \n",
    "        if self._return_attention:\n",
    "            attention = Concatenate()(attentions)\n",
    "            # print(\"attention\", attention, attention.shape)\n",
    "       \n",
    "        if self._return_attention:\n",
    "            return [output, attention]\n",
    "        else:\n",
    "            return output        \n",
    "\n",
    "# https://wanasit.github.io/attention-based-sequence-to-sequence-in-keras.html\n",
    "# https://arxiv.org/pdf/1508.04025.pdf\n",
    "class SequenceAttention(Layer):\n",
    "    \"\"\"\n",
    "        Takes two inputs of the shape (batch_size, T, dim1) and (batch_size, T, dim2),\n",
    "        whereby the first item is the source data and the second one the key data.\n",
    "        This layer then calculates for each batch's element and each time step a softmax attention \n",
    "        vector between the key data and the source data. Finally, this attention vector is multiplied\n",
    "        with the source data to obtain a weighted output. This means, that the key data is used to\n",
    "        interpret the source data in a special way to create an output of the same shape as the source data.\n",
    "    \"\"\"\n",
    "    def __init__(self, similarity, kernel_initializer=\"glorot_uniform\", **kwargs):\n",
    "        super(SequenceAttention, self).__init__(**kwargs)\n",
    "        if isinstance(similarity, str):\n",
    "            ALLOWED_SIMILARITIES = [\"additive\", \"multiplicative\" ]\n",
    "            if similarity not in ALLOWED_SIMILARITIES:\n",
    "                raise ValueError(\"`similarity` has to be either a callable or one of the following: {0}\".format(ALLOWED_SIMILARITIES))\n",
    "            else:\n",
    "                self._similarity = getattr(self, \"_\" + similarity + \"_similarity\")\n",
    "        elif callable(similarity):\n",
    "            self._similarity = similarity\n",
    "        else:\n",
    "            raise ValueError(\"`similarity` has to be either a callable or one of the following: {0}\".format(ALLOWED_SIMILARITIES))\n",
    "            \n",
    "        self._kernel_initializer = kernel_initializer\n",
    "            \n",
    "    def build(self, input_shape):\n",
    "        super(SequenceAttention, self).build(input_shape)\n",
    "        self._validate_input_shape(input_shape)\n",
    "        \n",
    "        self._weights = {}\n",
    "        if self._similarity == self._additive_similarity:\n",
    "            self._weights[\"w_a\"] = self.add_weight(\n",
    "                name='w_a', \n",
    "                shape=(input_shape[0][-1] + input_shape[1][-1], input_shape[0][-1]),\n",
    "                initializer=self._kernel_initializer,\n",
    "                trainable=True\n",
    "            )\n",
    "            \n",
    "            self._weights[\"v_a\"] = self.add_weight(\n",
    "                name='v_a', \n",
    "                shape=(1, input_shape[0][-1]),\n",
    "                initializer=self._kernel_initializer,\n",
    "                trainable=True\n",
    "            )\n",
    "            \n",
    "        elif self._similarity == self._multiplicative_similarity:\n",
    "            self._weights[\"w_a\"] = self.add_weight(\n",
    "                name='w_a', \n",
    "                shape=(input_shape[1][-1], input_shape[0][-1]),\n",
    "                initializer=self._kernel_initializer,\n",
    "                trainable=True\n",
    "            )\n",
    "\n",
    "        self.built = True\n",
    "        \n",
    "    def compute_output_shape(self, input_shape):\n",
    "        self._validate_input_shape(input_shape)\n",
    "        \n",
    "        return input_shape[0]\n",
    "            \n",
    "    def _validate_input_shape(self, input_shape):\n",
    "        if len(input_shape) != 2:\n",
    "            raise ValueError(\"Layer received an input shape {0} but expected two inputs (source, query).\".format(input_shape))\n",
    "        else:\n",
    "            if input_shape[0][0] != input_shape[1][0]:\n",
    "                raise ValueError(\"Both two inputs (source, query) have to have the same batch size; received batch sizes: {0}, {1}\".format(input_shape[0][0], input_shape[1][0]))\n",
    "            if input_shape[0][1] != input_shape[1][1]:\n",
    "                raise ValueError(\"Both inputs (source, query) have to have the same length; received lengths: {0}, {1}\".format(input_shape[0][0], input_shape[1][0]))\n",
    "        \n",
    "    def call(self, x):\n",
    "        source, query = x\n",
    "        \n",
    "        similarity = self._similarity(source, query)\n",
    "        expected_similarity_shape = [source.shape.as_list()[0], source.shape.as_list()[1], source.shape.as_list()[1]]\n",
    "       \n",
    "        if similarity.shape.as_list() != expected_similarity_shape:\n",
    "            raise RuntimeError(\"The similarity function has returned a similarity with shape {0}, but expected {1}\".format(similarity.shape.as_list()[:2], expected_similarity_shape))\n",
    "        \n",
    "        score = K.softmax(similarity)\n",
    "        output = K.batch_dot(score, source, axes=[1, 1])\n",
    "        \n",
    "        return output\n",
    "    \n",
    "    def _additive_similarity(self, source, query):\n",
    "        concatenation = K.concatenate([source, query], axis=2)\n",
    "        nonlinearity = K.tanh(K.dot(concatenation, self._weights[\"w_a\"]))\n",
    "        \n",
    "        # tile the weight vector (1, 1, dim) for each time step and each element of the batch -> (bs, T, dim)\n",
    "        source_shape = K.shape(source)\n",
    "        vaeff = K.tile(K.expand_dims(self._weights[\"v_a\"], 0), [source_shape[0], source_shape[1], 1])\n",
    "\n",
    "        similarity = K.batch_dot(K.permute_dimensions(vaeff, [0, 2, 1]), nonlinearity, axes=[1, 2])\n",
    "        \n",
    "        return similarity\n",
    "\n",
    "    def _multiplicative_similarity(self, source, query):\n",
    "        qp = K.dot(query, self._weights[\"w_a\"])\n",
    "        similarity = K.batch_dot(K.permute_dimensions(qp, [0, 2, 1]), source, axes=[1, 2])\n",
    "        \n",
    "        return similarity\n",
    "\n",
    "    def get_config(self):\n",
    "        config = {'similarity': self._similarity, 'kernel_initializer': self._kernel_initializer}\n",
    "        base_config = super(SequenceAttention, self).get_config()\n",
    "        return dict(list(base_config.items()) + list(config.items()))\n",
    "\n",
    "class AttentionRNNWrapper(Wrapper):\n",
    "    \"\"\"\n",
    "        The idea of the implementation is based on the paper:\n",
    "            \"Effective Approaches to Attention-based Neural Machine Translation\" by Luong et al.\n",
    "        This layer is an attention layer, which can be wrapped around arbitrary RNN layers.\n",
    "        This way, after each time step an attention vector is calculated\n",
    "        based on the current output of the LSTM and the entire input time series.\n",
    "        This attention vector is then used as a weight vector to choose special values\n",
    "        from the input data. This data is then finally concatenated to the next input\n",
    "        time step's data. On this a linear transformation in the same space as the input data's space\n",
    "        is performed before the data is fed into the RNN cell again.\n",
    "        This technique is similar to the input-feeding method described in the paper cited\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, layer, weight_initializer=\"glorot_uniform\", **kwargs):\n",
    "        assert isinstance(layer, RNN)\n",
    "        self.layer = layer\n",
    "        self.supports_masking = True\n",
    "        self.weight_initializer = weight_initializer\n",
    "        \n",
    "        super(AttentionRNNWrapper, self).__init__(layer, **kwargs)\n",
    "        \n",
    "    def _validate_input_shape(self, input_shape):\n",
    "        if len(input_shape) != 3:\n",
    "            raise ValueError(\"Layer received an input with shape {0} but expected a Tensor of rank 3.\".format(input_shape[0]))\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        self._validate_input_shape(input_shape)\n",
    "\n",
    "        self.input_spec = InputSpec(shape=input_shape)\n",
    "        \n",
    "        if not self.layer.built:\n",
    "            self.layer.build(input_shape)\n",
    "            self.layer.built = True\n",
    "            \n",
    "        input_dim = input_shape[-1]\n",
    "\n",
    "        if self.layer.return_sequences:\n",
    "            output_dim = self.layer.compute_output_shape(input_shape)[0][-1]\n",
    "        else:\n",
    "            output_dim = self.layer.compute_output_shape(input_shape)[-1]\n",
    "      \n",
    "        self._W1 = self.add_weight(shape=(input_dim, input_dim), name=\"{}_W1\".format(self.name), initializer=self.weight_initializer)\n",
    "        self._W2 = self.add_weight(shape=(output_dim, input_dim), name=\"{}_W2\".format(self.name), initializer=self.weight_initializer)\n",
    "        self._W3 = self.add_weight(shape=(2*input_dim, input_dim), name=\"{}_W3\".format(self.name), initializer=self.weight_initializer)\n",
    "        self._b2 = self.add_weight(shape=(input_dim,), name=\"{}_b2\".format(self.name), initializer=self.weight_initializer)\n",
    "        self._b3 = self.add_weight(shape=(input_dim,), name=\"{}_b3\".format(self.name), initializer=self.weight_initializer)\n",
    "        self._V = self.add_weight(shape=(input_dim,1), name=\"{}_V\".format(self.name), initializer=self.weight_initializer)\n",
    "        \n",
    "        super(AttentionRNNWrapper, self).build()\n",
    "        \n",
    "    def compute_output_shape(self, input_shape):\n",
    "        self._validate_input_shape(input_shape)\n",
    "\n",
    "        return self.layer.compute_output_shape(input_shape)\n",
    "    \n",
    "    @property\n",
    "    def trainable_weights(self):\n",
    "        return self._trainable_weights + self.layer.trainable_weights\n",
    "\n",
    "    @property\n",
    "    def non_trainable_weights(self):\n",
    "        return self._non_trainable_weights + self.layer.non_trainable_weights\n",
    "\n",
    "    def step(self, x, states):   \n",
    "        h = states[0]\n",
    "        # states[1] necessary?\n",
    "\n",
    "        # equals K.dot(X, self._W1) + self._b2 with X.shape=[bs, T, input_dim]\n",
    "        total_x_prod = states[-1]\n",
    "        # comes from the constants (equals the input sequence)\n",
    "        X = states[-2]\n",
    "        \n",
    "        # expand dims to add the vector which is only valid for this time step\n",
    "        # to total_x_prod which is valid for all time steps\n",
    "        hw = K.expand_dims(K.dot(h, self._W2), 1)\n",
    "        additive_atn = total_x_prod + hw\n",
    "        attention = K.softmax(K.dot(additive_atn, self._V), axis=1)\n",
    "        x_weighted = K.sum(attention * X, [1])\n",
    "\n",
    "        x = K.dot(K.concatenate([x, x_weighted], 1), self._W3) + self._b3\n",
    "        \n",
    "        h, new_states = self.layer.cell.call(x, states[:-2])\n",
    "        \n",
    "        return h, new_states\n",
    "    \n",
    "    def call(self, x, constants=None, mask=None, initial_state=None):\n",
    "        # input shape: (n_samples, time (padded with zeros), input_dim)\n",
    "        input_shape = self.input_spec.shape\n",
    "\n",
    "        if self.layer.stateful:\n",
    "            initial_states = self.layer.states\n",
    "        elif initial_state is not None:\n",
    "            initial_states = initial_state\n",
    "            if not isinstance(initial_states, (list, tuple)):\n",
    "                initial_states = [initial_states]\n",
    "\n",
    "            base_initial_state = self.layer.get_initial_state(x)\n",
    "            if len(base_initial_state) != len(initial_states):\n",
    "                raise ValueError(\"initial_state does not have the correct length. Received length {0} but expected {1}\".format(len(initial_states), len(base_initial_state)))\n",
    "            else:\n",
    "                # check the state' shape\n",
    "                for i in range(len(initial_states)):\n",
    "                    if not initial_states[i].shape.is_compatible_with(base_initial_state[i].shape): #initial_states[i][j] != base_initial_state[i][j]:\n",
    "                        raise ValueError(\"initial_state does not match the default base state of the layer. Received {0} but expected {1}\".format([x.shape for x in initial_states], [x.shape for x in base_initial_state]))\n",
    "        else:\n",
    "            initial_states = self.layer.get_initial_state(x)\n",
    "            \n",
    "        if not constants:\n",
    "            constants = []\n",
    "            \n",
    "        constants += self.get_constants(x)\n",
    "        \n",
    "        last_output, outputs, states = K.rnn(\n",
    "            self.step,\n",
    "            x,\n",
    "            initial_states,\n",
    "            go_backwards=self.layer.go_backwards,\n",
    "            mask=mask,\n",
    "            constants=constants,\n",
    "            unroll=self.layer.unroll,\n",
    "            input_length=input_shape[1]\n",
    "        )\n",
    "        \n",
    "        if self.layer.stateful:\n",
    "            self.updates = []\n",
    "            for i in range(len(states)):\n",
    "                self.updates.append((self.layer.states[i], states[i]))\n",
    "\n",
    "        if self.layer.return_sequences:\n",
    "            output = outputs\n",
    "        else:\n",
    "            output = last_output \n",
    "\n",
    "        # Properly set learning phase\n",
    "        if getattr(last_output, '_uses_learning_phase', False):\n",
    "            output._uses_learning_phase = True\n",
    "            for state in states:\n",
    "                state._uses_learning_phase = True\n",
    "\n",
    "        if self.layer.return_state:\n",
    "            if not isinstance(states, (list, tuple)):\n",
    "                states = [states]\n",
    "            else:\n",
    "                states = list(states)\n",
    "            return [output] + states\n",
    "        else:\n",
    "            return output\n",
    "\n",
    "    def get_constants(self, x):\n",
    "        # add constants to speed up calculation\n",
    "        constants = [x, K.dot(x, self._W1) + self._b2]\n",
    "        \n",
    "        return constants\n",
    "\n",
    "    def get_config(self):\n",
    "        config = {'weight_initializer': self.weight_initializer}\n",
    "        base_config = super(AttentionRNNWrapper, self).get_config()\n",
    "        return dict(list(base_config.items()) + list(config.items()))\n",
    "\n",
    "class ExternalAttentionRNNWrapper(Wrapper):\n",
    "    \"\"\"\n",
    "        The basic idea of the implementation is based on the paper:\n",
    "            \"Effective Approaches to Attention-based Neural Machine Translation\" by Luong et al.\n",
    "        This layer is an attention layer, which can be wrapped around arbitrary RNN layers.\n",
    "        This way, after each time step an attention vector is calculated\n",
    "        based on the current output of the LSTM and the entire input time series.\n",
    "        This attention vector is then used as a weight vector to choose special values\n",
    "        from the input data. This data is then finally concatenated to the next input\n",
    "        time step's data. On this a linear transformation in the same space as the input data's space\n",
    "        is performed before the data is fed into the RNN cell again.\n",
    "        This technique is similar to the input-feeding method described in the paper cited.\n",
    "        The only difference compared to the AttentionRNNWrapper is, that this layer\n",
    "        applies the attention layer not on the time-depending input but on a second\n",
    "        time-independent input (like image clues) as described in:\n",
    "            Show, Attend and Tell: Neural Image Caption Generation with Visual Attention\n",
    "            https://arxiv.org/abs/1502.03044\n",
    "    \"\"\"\n",
    "    def __init__(self, layer, weight_initializer=\"glorot_uniform\", return_attention=False, **kwargs):\n",
    "        assert isinstance(layer, RNN)\n",
    "        self.layer = layer\n",
    "        self.supports_masking = True\n",
    "        self.weight_initializer = weight_initializer\n",
    "        self.return_attention = return_attention\n",
    "        self._num_constants = None\n",
    "\n",
    "        super(ExternalAttentionRNNWrapper, self).__init__(layer, **kwargs)\n",
    "\n",
    "        self.input_spec = [InputSpec(ndim=3), InputSpec(ndim=3)]\n",
    "        \n",
    "    def _validate_input_shape(self, input_shape):\n",
    "        if len(input_shape) >= 2:\n",
    "            if len(input_shape[:2]) != 2:\n",
    "                raise ValueError(\"Layer has to receive two inputs: the temporal signal and the external signal which is constant for all time steps\")\n",
    "            if len(input_shape[0]) != 3:\n",
    "                raise ValueError(\"Layer received a temporal input with shape {0} but expected a Tensor of rank 3.\".format(input_shape[0]))\n",
    "            if len(input_shape[1]) != 3:\n",
    "                raise ValueError(\"Layer received a time-independent input with shape {0} but expected a Tensor of rank 3.\".format(input_shape[1]))\n",
    "        else:\n",
    "            raise ValueError(\"Layer has to receive at least 2 inputs: the temporal signal and the external signal which is constant for all time steps\")\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        self._validate_input_shape(input_shape)\n",
    "\n",
    "        for i, x in enumerate(input_shape):\n",
    "            self.input_spec[i] = InputSpec(shape=x)\n",
    "        \n",
    "        if not self.layer.built:\n",
    "            self.layer.build(input_shape)\n",
    "            self.layer.built = True\n",
    "            \n",
    "        temporal_input_dim = input_shape[0][-1]\n",
    "        static_input_dim = input_shape[1][-1]\n",
    "\n",
    "        if self.layer.return_sequences:\n",
    "            output_dim = self.layer.compute_output_shape(input_shape[0])[0][-1]\n",
    "        else:\n",
    "            output_dim = self.layer.compute_output_shape(input_shape[0])[-1]\n",
    "      \n",
    "        self._W1 = self.add_weight(shape=(static_input_dim, temporal_input_dim), name=\"{}_W1\".format(self.name), initializer=self.weight_initializer)\n",
    "        self._W2 = self.add_weight(shape=(output_dim, temporal_input_dim), name=\"{}_W2\".format(self.name), initializer=self.weight_initializer)\n",
    "        self._W3 = self.add_weight(shape=(temporal_input_dim + static_input_dim, temporal_input_dim), name=\"{}_W3\".format(self.name), initializer=self.weight_initializer)\n",
    "        self._b2 = self.add_weight(shape=(temporal_input_dim,), name=\"{}_b2\".format(self.name), initializer=self.weight_initializer)\n",
    "        self._b3 = self.add_weight(shape=(temporal_input_dim,), name=\"{}_b3\".format(self.name), initializer=self.weight_initializer)\n",
    "        self._V = self.add_weight(shape=(temporal_input_dim, 1), name=\"{}_V\".format(self.name), initializer=self.weight_initializer)\n",
    "        \n",
    "        super(ExternalAttentionRNNWrapper, self).build()\n",
    "        \n",
    "    @property\n",
    "    def trainable_weights(self):\n",
    "        return self._trainable_weights + self.layer.trainable_weights\n",
    "\n",
    "    @property\n",
    "    def non_trainable_weights(self):\n",
    "        return self._non_trainable_weights + self.layer.non_trainable_weights\n",
    "\n",
    "    def compute_output_shape(self, input_shape):\n",
    "        self._validate_input_shape(input_shape)\n",
    "\n",
    "        output_shape =  self.layer.compute_output_shape(input_shape[0])\n",
    "\n",
    "        if self.return_attention:\n",
    "            if not isinstance(output_shape, list):\n",
    "                output_shape = [output_shape]\n",
    "\n",
    "            output_shape = output_shape + [(None, input_shape[1][1])]\n",
    "\n",
    "        return output_shape\n",
    "    \n",
    "    def step(self, x, states):  \n",
    "        h = states[0]\n",
    "        # states[1] necessary?\n",
    "        \n",
    "        # comes from the constants\n",
    "        X_static = states[-2]\n",
    "        # equals K.dot(static_x, self._W1) + self._b2 with X.shape=[bs, L, static_input_dim]\n",
    "        total_x_static_prod = states[-1]\n",
    "\n",
    "        # expand dims to add the vector which is only valid for this time step\n",
    "        # to total_x_prod which is valid for all time steps\n",
    "        hw = K.expand_dims(K.dot(h, self._W2), 1)\n",
    "        additive_atn = total_x_static_prod + hw\n",
    "        attention = K.softmax(K.dot(additive_atn, self._V), axis=1)\n",
    "        static_x_weighted = K.sum(attention * X_static, [1])\n",
    "        \n",
    "        x = K.dot(K.concatenate([x, static_x_weighted], 1), self._W3) + self._b3\n",
    "\n",
    "        h, new_states = self.layer.cell.call(x, states[:-2])\n",
    "        \n",
    "        # append attention to the states to \"smuggle\" it out of the RNN wrapper\n",
    "        attention = K.squeeze(attention, -1)\n",
    "        h = K.concatenate([h, attention])\n",
    "\n",
    "        return h, new_states\n",
    "    \n",
    "    def call(self, x, constants=None, mask=None, initial_state=None):\n",
    "        # input shape: (n_samples, time (padded with zeros), input_dim)\n",
    "        input_shape = self.input_spec[0].shape\n",
    "\n",
    "        if len(x) > 2:\n",
    "            initial_state = x[2:]\n",
    "            x = x[:2]\n",
    "            assert len(initial_state) >= 1\n",
    "\n",
    "        static_x = x[1]\n",
    "        x = x[0]\n",
    "\n",
    "        if self.layer.stateful:\n",
    "            initial_states = self.layer.states\n",
    "        elif initial_state is not None:\n",
    "            initial_states = initial_state\n",
    "            if not isinstance(initial_states, (list, tuple)):\n",
    "                initial_states = [initial_states]\n",
    "        else:\n",
    "            initial_states = self.layer.get_initial_state(x)\n",
    "            \n",
    "        if not constants:\n",
    "            constants = []\n",
    "        constants += self.get_constants(static_x)\n",
    "\n",
    "        last_output, outputs, states = K.rnn(\n",
    "            self.step,\n",
    "            x,\n",
    "            initial_states,\n",
    "            go_backwards=self.layer.go_backwards,\n",
    "            mask=mask,\n",
    "            constants=constants,\n",
    "            unroll=self.layer.unroll,\n",
    "            input_length=input_shape[1]\n",
    "        )\n",
    "\n",
    "        # output has at the moment the form:\n",
    "        # (real_output, attention)\n",
    "        # split this now up\n",
    "\n",
    "        output_dim = self.layer.compute_output_shape(input_shape)[0][-1]\n",
    "        last_output = last_output[:output_dim]\n",
    "\n",
    "        attentions = outputs[:, :, output_dim:]\n",
    "        outputs = outputs[:, :, :output_dim]\n",
    "        \n",
    "        if self.layer.stateful:\n",
    "            self.updates = []\n",
    "            for i in range(len(states)):\n",
    "                self.updates.append((self.layer.states[i], states[i]))\n",
    "\n",
    "        if self.layer.return_sequences:\n",
    "            output = outputs\n",
    "        else:\n",
    "            output = last_output \n",
    "\n",
    "        # Properly set learning phase\n",
    "        if getattr(last_output, '_uses_learning_phase', False):\n",
    "            output._uses_learning_phase = True\n",
    "            for state in states:\n",
    "                state._uses_learning_phase = True\n",
    "\n",
    "        if self.layer.return_state:\n",
    "            if not isinstance(states, (list, tuple)):\n",
    "                states = [states]\n",
    "            else:\n",
    "                states = list(states)\n",
    "            output = [output] + states\n",
    "\n",
    "        if self.return_attention:\n",
    "            if not isinstance(output, list):\n",
    "                output = [output]\n",
    "            output = output + [attentions]\n",
    "\n",
    "        return output\n",
    "\n",
    "    def _standardize_args(self, inputs, initial_state, constants, num_constants):\n",
    "        \"\"\"Standardize `__call__` to a single list of tensor inputs.\n",
    "        When running a model loaded from file, the input tensors\n",
    "        `initial_state` and `constants` can be passed to `RNN.__call__` as part\n",
    "        of `inputs` instead of by the dedicated keyword arguments. This method\n",
    "        makes sure the arguments are separated and that `initial_state` and\n",
    "        `constants` are lists of tensors (or None).\n",
    "        # Arguments\n",
    "        inputs: tensor or list/tuple of tensors\n",
    "        initial_state: tensor or list of tensors or None\n",
    "        constants: tensor or list of tensors or None\n",
    "        # Returns\n",
    "        inputs: tensor\n",
    "        initial_state: list of tensors or None\n",
    "        constants: list of tensors or None\n",
    "        \"\"\"\n",
    "        inputs=inputs[:2]\n",
    "        if isinstance(inputs, list) and len(inputs) > 2:\n",
    "            assert initial_state is None and constants is None\n",
    "            if num_constants is not None:\n",
    "                constants = inputs[-num_constants:]\n",
    "                inputs = inputs[:-num_constants]\n",
    "            initial_state = inputs[2:]\n",
    "            inputs = inputs[:2]\n",
    "\n",
    "        def to_list_or_none(x):\n",
    "            if x is None or isinstance(x, list):\n",
    "                return x\n",
    "            if isinstance(x, tuple):\n",
    "                return list(x)\n",
    "            return [x]\n",
    "\n",
    "        initial_state = to_list_or_none(initial_state)\n",
    "        constants = to_list_or_none(constants)\n",
    "\n",
    "        return inputs, initial_state, constants\n",
    "\n",
    "    def __call__(self, inputs, initial_state=None, constants=None, **kwargs):\n",
    "        inputs, initial_state, constants = self._standardize_args(\n",
    "            inputs, initial_state, constants, self._num_constants)\n",
    "\n",
    "        if initial_state is None and constants is None:\n",
    "            return super(ExternalAttentionRNNWrapper, self).__call__(inputs, **kwargs)\n",
    "\n",
    "        # If any of `initial_state` or `constants` are specified and are Keras\n",
    "        # tensors, then add them to the inputs and temporarily modify the\n",
    "        # input_spec to include them.\n",
    "\n",
    "        additional_inputs = []\n",
    "        additional_specs = []\n",
    "        if initial_state is not None:\n",
    "            kwargs['initial_state'] = initial_state\n",
    "            additional_inputs += initial_state\n",
    "            self.state_spec = [InputSpec(shape=K.int_shape(state))\n",
    "                               for state in initial_state]\n",
    "            additional_specs += self.state_spec\n",
    "        if constants is not None:\n",
    "            kwargs['constants'] = constants\n",
    "            additional_inputs += constants\n",
    "            self.constants_spec = [InputSpec(shape=K.int_shape(constant))\n",
    "                                   for constant in constants]\n",
    "            self._num_constants = len(constants)\n",
    "            additional_specs += self.constants_spec\n",
    "        # at this point additional_inputs cannot be empty\n",
    "        is_keras_tensor = K.is_keras_tensor(additional_inputs[0])\n",
    "        for tensor in additional_inputs:\n",
    "            if K.is_keras_tensor(tensor) != is_keras_tensor:\n",
    "                raise ValueError('The initial state or constants of an ExternalAttentionRNNWrapper'\n",
    "                                 ' layer cannot be specified with a mix of'\n",
    "                                 ' Keras tensors and non-Keras tensors'\n",
    "                                 ' (a \"Keras tensor\" is a tensor that was'\n",
    "                                 ' returned by a Keras layer, or by `Input`)')\n",
    "\n",
    "        if is_keras_tensor:\n",
    "            # Compute the full input spec, including state and constants\n",
    "            full_input = inputs + additional_inputs\n",
    "            full_input_spec = self.input_spec + additional_specs\n",
    "            # Perform the call with temporarily replaced input_spec\n",
    "            original_input_spec = self.input_spec\n",
    "            self.input_spec = full_input_spec\n",
    "            output = super(ExternalAttentionRNNWrapper, self).__call__(full_input, **kwargs)\n",
    "            self.input_spec = self.input_spec[:len(original_input_spec)]\n",
    "            return output\n",
    "        else:\n",
    "            return super(ExternalAttentionRNNWrapper, self).__call__(inputs, **kwargs)\n",
    "\n",
    "    def get_constants(self, x):\n",
    "        # add constants to speed up calculation\n",
    "        constants = [x, K.dot(x, self._W1) + self._b2]\n",
    "        return constants\n",
    "\n",
    "    def get_config(self):\n",
    "        config = {'return_attention': self.return_attention, 'weight_initializer': self.weight_initializer}\n",
    "        base_config = super(ExternalAttentionRNNWrapper, self).get_config()\n",
    "        return dict(list(base_config.items()) + list(config.items()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import keras\n",
    "\n",
    "\n",
    "\n",
    "class SeqSelfAttention(keras.layers.Layer):\n",
    "\n",
    "    ATTENTION_TYPE_ADD = 'additive'\n",
    "    ATTENTION_TYPE_MUL = 'multiplicative'\n",
    "\n",
    "    def __init__(self,\n",
    "                 units=32,\n",
    "                 attention_width=None,\n",
    "                 attention_type=ATTENTION_TYPE_ADD,\n",
    "                 return_attention=False,\n",
    "                 history_only=False,\n",
    "                 kernel_initializer='glorot_normal',\n",
    "                 bias_initializer='zeros',\n",
    "                 kernel_regularizer=None,\n",
    "                 bias_regularizer=None,\n",
    "                 kernel_constraint=None,\n",
    "                 bias_constraint=None,\n",
    "                 use_additive_bias=True,\n",
    "                 use_attention_bias=True,\n",
    "                 attention_activation=None,\n",
    "                 attention_regularizer_weight=0.0,\n",
    "                 **kwargs):\n",
    "        \"\"\"Layer initialization.\n",
    "        For additive attention, see: https://arxiv.org/pdf/1806.01264.pdf\n",
    "        :param units: The dimension of the vectors that used to calculate the attention weights.\n",
    "        :param attention_width: The width of local attention.\n",
    "        :param attention_type: 'additive' or 'multiplicative'.\n",
    "        :param return_attention: Whether to return the attention weights for visualization.\n",
    "        :param history_only: Only use historical pieces of data.\n",
    "        :param kernel_initializer: The initializer for weight matrices.\n",
    "        :param bias_initializer: The initializer for biases.\n",
    "        :param kernel_regularizer: The regularization for weight matrices.\n",
    "        :param bias_regularizer: The regularization for biases.\n",
    "        :param kernel_constraint: The constraint for weight matrices.\n",
    "        :param bias_constraint: The constraint for biases.\n",
    "        :param use_additive_bias: Whether to use bias while calculating the relevance of inputs features\n",
    "                                  in additive mode.\n",
    "        :param use_attention_bias: Whether to use bias while calculating the weights of attention.\n",
    "        :param attention_activation: The activation used for calculating the weights of attention.\n",
    "        :param attention_regularizer_weight: The weights of attention regularizer.\n",
    "        :param kwargs: Parameters for parent class.\n",
    "        \"\"\"\n",
    "        super(SeqSelfAttention, self).__init__(**kwargs)\n",
    "        self.supports_masking = True\n",
    "        self.units = units\n",
    "        self.attention_width = attention_width\n",
    "        self.attention_type = attention_type\n",
    "        self.return_attention = return_attention\n",
    "        self.history_only = history_only\n",
    "        if history_only and attention_width is None:\n",
    "            self.attention_width = int(1e9)\n",
    "\n",
    "        self.use_additive_bias = use_additive_bias\n",
    "        self.use_attention_bias = use_attention_bias\n",
    "        self.kernel_initializer = keras.initializers.get(kernel_initializer)\n",
    "        self.bias_initializer = keras.initializers.get(bias_initializer)\n",
    "        self.kernel_regularizer = keras.regularizers.get(kernel_regularizer)\n",
    "        self.bias_regularizer = keras.regularizers.get(bias_regularizer)\n",
    "        self.kernel_constraint = keras.constraints.get(kernel_constraint)\n",
    "        self.bias_constraint = keras.constraints.get(bias_constraint)\n",
    "        self.attention_activation = keras.activations.get(attention_activation)\n",
    "        self.attention_regularizer_weight = attention_regularizer_weight\n",
    "        self._backend = keras.backend.backend()\n",
    "\n",
    "        if attention_type == SeqSelfAttention.ATTENTION_TYPE_ADD:\n",
    "            self.Wx, self.Wt, self.bh = None, None, None\n",
    "            self.Wa, self.ba = None, None\n",
    "        elif attention_type == SeqSelfAttention.ATTENTION_TYPE_MUL:\n",
    "            self.Wa, self.ba = None, None\n",
    "        else:\n",
    "            raise NotImplementedError('No implementation for attention type : ' + attention_type)\n",
    "\n",
    "    def get_config(self):\n",
    "        config = {\n",
    "            'units': self.units,\n",
    "            'attention_width': self.attention_width,\n",
    "            'attention_type': self.attention_type,\n",
    "            'return_attention': self.return_attention,\n",
    "            'history_only': self.history_only,\n",
    "            'use_additive_bias': self.use_additive_bias,\n",
    "            'use_attention_bias': self.use_attention_bias,\n",
    "            'kernel_initializer': keras.initializers.serialize(self.kernel_initializer),\n",
    "            'bias_initializer': keras.initializers.serialize(self.bias_initializer),\n",
    "            'kernel_regularizer': keras.regularizers.serialize(self.kernel_regularizer),\n",
    "            'bias_regularizer': keras.regularizers.serialize(self.bias_regularizer),\n",
    "            'kernel_constraint': keras.constraints.serialize(self.kernel_constraint),\n",
    "            'bias_constraint': keras.constraints.serialize(self.bias_constraint),\n",
    "            'attention_activation': keras.activations.serialize(self.attention_activation),\n",
    "            'attention_regularizer_weight': self.attention_regularizer_weight,\n",
    "        }\n",
    "        base_config = super(SeqSelfAttention, self).get_config()\n",
    "        return dict(list(base_config.items()) + list(config.items()))\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        if self.attention_type == SeqSelfAttention.ATTENTION_TYPE_ADD:\n",
    "            self._build_additive_attention(input_shape)\n",
    "        elif self.attention_type == SeqSelfAttention.ATTENTION_TYPE_MUL:\n",
    "            self._build_multiplicative_attention(input_shape)\n",
    "        super(SeqSelfAttention, self).build(input_shape)\n",
    "\n",
    "    def _build_additive_attention(self, input_shape):\n",
    "        feature_dim = int(input_shape[2])\n",
    "\n",
    "        self.Wt = self.add_weight(shape=(feature_dim, self.units),\n",
    "                                  name='{}_Add_Wt'.format(self.name),\n",
    "                                  initializer=self.kernel_initializer,\n",
    "                                  regularizer=self.kernel_regularizer,\n",
    "                                  constraint=self.kernel_constraint)\n",
    "        self.Wx = self.add_weight(shape=(feature_dim, self.units),\n",
    "                                  name='{}_Add_Wx'.format(self.name),\n",
    "                                  initializer=self.kernel_initializer,\n",
    "                                  regularizer=self.kernel_regularizer,\n",
    "                                  constraint=self.kernel_constraint)\n",
    "        if self.use_additive_bias:\n",
    "            self.bh = self.add_weight(shape=(self.units,),\n",
    "                                      name='{}_Add_bh'.format(self.name),\n",
    "                                      initializer=self.bias_initializer,\n",
    "                                      regularizer=self.bias_regularizer,\n",
    "                                      constraint=self.bias_constraint)\n",
    "\n",
    "        self.Wa = self.add_weight(shape=(self.units, 1),\n",
    "                                  name='{}_Add_Wa'.format(self.name),\n",
    "                                  initializer=self.kernel_initializer,\n",
    "                                  regularizer=self.kernel_regularizer,\n",
    "                                  constraint=self.kernel_constraint)\n",
    "        if self.use_attention_bias:\n",
    "            self.ba = self.add_weight(shape=(1,),\n",
    "                                      name='{}_Add_ba'.format(self.name),\n",
    "                                      initializer=self.bias_initializer,\n",
    "                                      regularizer=self.bias_regularizer,\n",
    "                                      constraint=self.bias_constraint)\n",
    "\n",
    "    def _build_multiplicative_attention(self, input_shape):\n",
    "        feature_dim = int(input_shape[2])\n",
    "\n",
    "        self.Wa = self.add_weight(shape=(feature_dim, feature_dim),\n",
    "                                  name='{}_Mul_Wa'.format(self.name),\n",
    "                                  initializer=self.kernel_initializer,\n",
    "                                  regularizer=self.kernel_regularizer,\n",
    "                                  constraint=self.kernel_constraint)\n",
    "        if self.use_attention_bias:\n",
    "            self.ba = self.add_weight(shape=(1,),\n",
    "                                      name='{}_Mul_ba'.format(self.name),\n",
    "                                      initializer=self.bias_initializer,\n",
    "                                      regularizer=self.bias_regularizer,\n",
    "                                      constraint=self.bias_constraint)\n",
    "\n",
    "    def call(self, inputs, mask=None, **kwargs):\n",
    "        input_len = K.shape(inputs)[1]\n",
    "\n",
    "        if self.attention_type == SeqSelfAttention.ATTENTION_TYPE_ADD:\n",
    "            e = self._call_additive_emission(inputs)\n",
    "        elif self.attention_type == SeqSelfAttention.ATTENTION_TYPE_MUL:\n",
    "            e = self._call_multiplicative_emission(inputs)\n",
    "\n",
    "        if self.attention_activation is not None:\n",
    "            e = self.attention_activation(e)\n",
    "        if self.attention_width is not None:\n",
    "            if self.history_only:\n",
    "                lower = K.arange(0, input_len) - (self.attention_width - 1)\n",
    "            else:\n",
    "                lower = K.arange(0, input_len) - self.attention_width // 2\n",
    "            lower = K.expand_dims(lower, axis=-1)\n",
    "            upper = lower + self.attention_width\n",
    "            indices = K.expand_dims(K.arange(0, input_len), axis=0)\n",
    "            e -= 10000.0 * (1.0 - K.cast(lower <= indices, K.floatx()) * K.cast(indices < upper, K.floatx()))\n",
    "        if mask is not None:\n",
    "            mask = K.expand_dims(K.cast(mask, K.floatx()), axis=-1)\n",
    "            e -= 10000.0 * ((1.0 - mask) * (1.0 - K.permute_dimensions(mask, (0, 2, 1))))\n",
    "\n",
    "        # a_{t} = \\text{softmax}(e_t)\n",
    "        e = K.exp(e - K.max(e, axis=-1, keepdims=True))\n",
    "        a = e / K.sum(e, axis=-1, keepdims=True)\n",
    "\n",
    "        # l_t = \\sum_{t'} a_{t, t'} x_{t'}\n",
    "        v = K.batch_dot(a, inputs)\n",
    "        if self.attention_regularizer_weight > 0.0:\n",
    "            self.add_loss(self._attention_regularizer(a))\n",
    "\n",
    "        if self.return_attention:\n",
    "            return [v, a]\n",
    "        return v\n",
    "\n",
    "    def _call_additive_emission(self, inputs):\n",
    "        input_shape = K.shape(inputs)\n",
    "        batch_size, input_len = input_shape[0], input_shape[1]\n",
    "\n",
    "        # h_{t, t'} = \\tanh(x_t^T W_t + x_{t'}^T W_x + b_h)\n",
    "        q = K.expand_dims(K.dot(inputs, self.Wt), 2)\n",
    "        k = K.expand_dims(K.dot(inputs, self.Wx), 1)\n",
    "        if self.use_additive_bias:\n",
    "            h = K.tanh(q + k + self.bh)\n",
    "        else:\n",
    "            h = K.tanh(q + k)\n",
    "\n",
    "        # e_{t, t'} = W_a h_{t, t'} + b_a\n",
    "        if self.use_attention_bias:\n",
    "            e = K.reshape(K.dot(h, self.Wa) + self.ba, (batch_size, input_len, input_len))\n",
    "        else:\n",
    "            e = K.reshape(K.dot(h, self.Wa), (batch_size, input_len, input_len))\n",
    "        return e\n",
    "\n",
    "    def _call_multiplicative_emission(self, inputs):\n",
    "        # e_{t, t'} = x_t^T W_a x_{t'} + b_a\n",
    "        e = K.batch_dot(K.dot(inputs, self.Wa), K.permute_dimensions(inputs, (0, 2, 1)))\n",
    "        if self.use_attention_bias:\n",
    "            e += self.ba[0]\n",
    "        return e\n",
    "\n",
    "    def compute_output_shape(self, input_shape):\n",
    "        output_shape = input_shape\n",
    "        if self.return_attention:\n",
    "            attention_shape = (input_shape[0], output_shape[1], input_shape[1])\n",
    "            return [output_shape, attention_shape]\n",
    "        return output_shape\n",
    "\n",
    "    def compute_mask(self, inputs, mask=None):\n",
    "        if self.return_attention:\n",
    "            return [mask, None]\n",
    "        return mask\n",
    "\n",
    "    def _attention_regularizer(self, attention):\n",
    "        batch_size = K.cast(K.shape(attention)[0], K.floatx())\n",
    "        input_len = K.shape(attention)[-1]\n",
    "        indices = K.expand_dims(K.arange(0, input_len), axis=0)\n",
    "        diagonal = K.expand_dims(K.arange(0, input_len), axis=-1)\n",
    "        eye = K.cast(K.equal(indices, diagonal), K.floatx())\n",
    "        return self.attention_regularizer_weight * K.sum(K.square(K.batch_dot(\n",
    "            attention,\n",
    "            K.permute_dimensions(attention, (0, 2, 1))) - eye)) / batch_size\n",
    "\n",
    "    @staticmethod\n",
    "    def get_custom_objects():\n",
    "        return {'SeqSelfAttention': SeqSelfAttention}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.Encoder import Encoder\n",
    "from utils.Classifier import Classifier\n",
    "from utils.Detector import Detector\n",
    "from utils.Segmenter import Segmenter\n",
    "\n",
    "\n",
    "img_size = 256\n",
    "n_classes = 1\n",
    "    \n",
    "encoder = Encoder(weights=None)\n",
    "classifier = Classifier(encoder)\n",
    "detector = Detector(encoder, img_size, n_classes)\n",
    "segmenter = Segmenter(encoder)\n",
    "heads=[]\n",
    "heads.append(classifier)\n",
    "heads.append(detector)\n",
    "heads.append(segmenter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#from kulc.attention import ExternalAttentionRNNWrapper\n",
    "from tensorflow.keras.layers import Input,Embedding,Lambda,Dense,TimeDistributed,LSTM,Reshape,Dropout\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras.models import Model\n",
    "import tensorflow.keras.backend as K\n",
    "from tensorflow.keras.utils import plot_model\n",
    "\n",
    "def create_model(encoder_model, vocabulary_size, embedding_size, T, L, D):\n",
    "    \n",
    "#     image_input = Input(shape=(256,256,3), name=\"image_input\")\n",
    "    \n",
    "#     image_model = tf.keras.applications.DenseNet121(include_top=False,\n",
    "#                                           input_shape=(256,256,3),\n",
    "#                                           weights='imagenet')\n",
    "    \n",
    "#     image_model = Model(image_model.input, image_model.layers[-2].output)\n",
    "    \n",
    "    #for layer in image_model.layers:\n",
    "    #    layer.trainable = False\n",
    "    \n",
    "    \n",
    "    image_features_input = encoder_model.model.output #image_model(image_input)\n",
    "    #image_features_input = Reshape((16*16,512))(image_features_input)\n",
    "    image_features_input = Reshape((8*8,1024))(image_features_input)\n",
    "    #image_features_input = Dropout(0.2)(image_features_input)\n",
    "    \n",
    "    captions_input = Input(shape=(T,), name=\"captions_input\")\n",
    "    captions = Embedding(vocabulary_size, embedding_size, input_length=T)(captions_input)\n",
    "\n",
    "    averaged_image_features = Lambda(lambda x: K.mean(x, axis=1))\n",
    "    averaged_image_features = averaged_image_features(image_features_input)\n",
    "    initial_state_h = Dense(embedding_size)(averaged_image_features)\n",
    "    initial_state_c = Dense(embedding_size)(averaged_image_features)\n",
    "  \n",
    "    image_features = TimeDistributed(Dense(D, activation=\"relu\" ))(image_features_input)\n",
    "    #image_features = Dropout(0.2)(image_features)\n",
    "    \n",
    "    encoder = LSTM(embedding_size, return_sequences=True, return_state=True, recurrent_dropout=0.1)\n",
    "    attented_encoder = ExternalAttentionRNNWrapper(encoder, return_attention=True )\n",
    "    self_attention_layer = SeqSelfAttention(attention_activation='relu')\n",
    "    \n",
    "    output = TimeDistributed(Dense(vocabulary_size, activation=\"softmax\"), name=\"output\")\n",
    "\n",
    "    # for training purpose\n",
    "    attented_encoder_training_data, _, _ , _= attented_encoder([captions, image_features], initial_state=[initial_state_h, initial_state_c])\n",
    "    \n",
    "    training_output_data = self_attention_layer(attented_encoder_training_data)\n",
    "    training_output_data = output(training_output_data)\n",
    "    \n",
    "    \n",
    "    \n",
    "    training_model = Model(inputs=[encoder_model.model.input,captions_input], outputs=[heads[0].model,heads[1].model,heads[2].model,training_output_data])\n",
    "    \n",
    "    initial_state_inference_model = Model(inputs=[encoder_model.model.input], outputs=[heads[0].model,heads[1].model,heads[2].model,image_features, initial_state_h, initial_state_c])\n",
    "    \n",
    "    inference_initial_state_h = Input(shape=(embedding_size,))\n",
    "    inference_initial_state_c = Input(shape=(embedding_size,))\n",
    "    image_input_feat = Input(shape=(64,D,))\n",
    "    \n",
    "    attented_encoder_inference_data, inference_encoder_state_h, inference_encoder_state_c, inference_attention = attented_encoder(\n",
    "        [captions, image_input_feat],\n",
    "        initial_state=[inference_initial_state_h, inference_initial_state_c]\n",
    "        )\n",
    "   \n",
    "    inference_output_data = self_attention_layer(attented_encoder_inference_data)\n",
    "    inference_output_data = output(inference_output_data)\n",
    "     \n",
    "    \n",
    "    \n",
    "    inference_model = Model(\n",
    "        inputs=[image_input_feat, captions_input, inference_initial_state_h, inference_initial_state_c],\n",
    "        outputs=[inference_output_data, inference_encoder_state_h, inference_encoder_state_c, inference_attention]\n",
    "    )\n",
    "    \n",
    "    return training_model, inference_model, initial_state_inference_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from data_loader.MTL_dataloader import get_train_validation_generator\n",
    "\n",
    "det_csv_path = \"/kaggle/input/rsna-pneumonia-detection-challenge/stage_2_train_labels.csv\"\n",
    "seg_csv_path = \"/kaggle/input/siim-acr-pneumothorax-segmentation-data/train-rle.csv\"\n",
    "det_images_path = \"/kaggle/input/rsna-pneumonia-detection-challenge/stage_2_train_images/\"\n",
    "seg_images_path = \"/kaggle/input/siim-acr-pneumothorax-segmentation-data/dicom-images-train/\"\n",
    "\n",
    "report_csv1_path = \"/kaggle/input/chest-xrays-indiana-university/indiana_reports.csv\"\n",
    "report_csv2_path = \"/kaggle/input/chest-xrays-indiana-university/indiana_projections.csv\"\n",
    "report_images_path=\"/kaggle/input/chest-xrays-indiana-university/images/images_normalized/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_gen,val_gen = get_train_validation_generator(det_csv_path,seg_csv_path , det_images_path, seg_images_path,\n",
    "                                                   report_csv1_path,\n",
    "                                                   report_csv2_path,\n",
    "                                                   report_images_path,\n",
    "                                                   augmentation=True,hist_eq=True,normalize=True ,only_positive=False,batch_positive_portion=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X,Y = next(enumerate(train_gen))[1]\n",
    "print(X[0].shape)\n",
    "print(X[1].shape)\n",
    "\n",
    "print(Y[0].shape)\n",
    "print(Y[1].shape)\n",
    "print(Y[2].shape)\n",
    "print(Y[3].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.Encoder import Encoder\n",
    "from utils.Classifier import Classifier\n",
    "from utils.Detector import Detector\n",
    "from utils.Segmenter import Segmenter\n",
    "\n",
    "img_size = 256\n",
    "n_classes = 1\n",
    "    \n",
    "encoder = Encoder(weights=None)\n",
    "classifier = Classifier(encoder)\n",
    "detector = Detector(encoder, img_size, n_classes)\n",
    "segmenter = Segmenter(encoder)\n",
    "heads=[]\n",
    "heads.append(classifier)\n",
    "heads.append(detector)\n",
    "heads.append(segmenter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_size = 512\n",
    "T= None\n",
    "L= 8*8\n",
    "D= 512\n",
    "\n",
    "vocab_size = train_gen.report_gen.vocab_size\n",
    "\n",
    "training_model, inference_model, initial_state_inference_model = create_model(encoder, vocab_size, embedding_size, T, L, D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_model.load_weights('../input/mtl-with-report-weights/7.hdf5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.losses import sparse_categorical_crossentropy\n",
    "import tensorflow as tf\n",
    "\n",
    "def loss_d(y_true,y_pred):\n",
    "    output =  tf.cond(\n",
    "                tf.math.reduce_all(tf.math.equal(y_true,-1))\n",
    "                ,true_fn= lambda: tf.convert_to_tensor(0, dtype=tf.float32)\n",
    "                ,false_fn= lambda: sparse_categorical_crossentropy(y_true,y_pred) )\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.losses import categorical_crossentropy\n",
    "\n",
    "def class_loss(y_true,y_pred):\n",
    "    return tf.cond(\n",
    "                    tf.math.reduce_all(tf.math.equal(y_true,-1))\n",
    "                    ,true_fn=  lambda: tf.convert_to_tensor(0, dtype=tf.float32)\n",
    "                    ,false_fn= lambda: categorical_crossentropy(y_true,y_pred)\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.optimizers import Adam\n",
    "epochs = 1\n",
    "lr=1e-4\n",
    "training_model.compile(loss=[classifier.loss , detector.loss , segmenter.loss , loss_d ],optimizer=Adam(lr))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " callback = tf.keras.callbacks.ModelCheckpoint(\n",
    "    filepath=\"./\",\n",
    "    save_weights_only=True,\n",
    "    monitor='val_loss',\n",
    "    mode='min',\n",
    "    save_freq=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this way the nan doesn't occur in training\n",
    "train_gen.nb_iteration = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for k in range(200):\n",
    "    training_model.fit_generator( train_gen,\n",
    "                    epochs = epochs,\n",
    "                    callbacks=[callback]\n",
    "                                   )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_model.save_weights(\"8.hdf5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def tokens_to_text(tokens,tok,end_token='endseq'):\n",
    "    sentence=\"\"\n",
    "    for token in tokens:\n",
    "        if token ==0:\n",
    "            break\n",
    "        \n",
    "        word = tok.index_word[token]\n",
    "        \n",
    "        if word==end_token:\n",
    "            break\n",
    "            \n",
    "        sentence+= word+\" \"\n",
    "        \n",
    "    sentence = sentence.strip()\n",
    "    \n",
    "    return sentence\n",
    "\n",
    "def predict(image,tok, initial_state_inference_model, inference_model,Y,start_token='startseq',end_token='endseq', max_len=100):\n",
    "\n",
    "    _,_,_,image_features, init_h,init_c = initial_state_inference_model(np.expand_dims(X[0][0],axis=0))\n",
    "    word = tok.word_index[start_token]\n",
    "\n",
    "    predictions=[]\n",
    "    for index in range(max_len):\n",
    "        #word = Y[index]\n",
    "        \n",
    "        word, init_h, init_c, inference_attention = inference_model([  np.array(image_features),\n",
    "                                                                       np.array([[word]]),\n",
    "                                                                       np.array(init_h),\n",
    "                                                                       np.array(init_c)  ] )\n",
    "        \n",
    "        \n",
    "        \n",
    "        word = tf.expand_dims(tf.squeeze(word),axis=0)\n",
    "        \n",
    "        word = tf.random.categorical(word, 1)[0][0].numpy()\n",
    "        \n",
    "        \n",
    "        \n",
    "        if word==tok.word_index[end_token]:\n",
    "            break\n",
    "\n",
    "        predictions.append(word)\n",
    "        \n",
    "    return predictions\n",
    "    \n",
    "    \n",
    "def get_sentence_preds(image,tok, initial_state_inference_model, inference_model,Y,start_token='startseq',end_token='endseq', max_len=100):\n",
    "    tokens = predict(image,tok,initial_state_inference_model,inference_model,Y)\n",
    "    sentence = tokens_to_text(tokens,tok)\n",
    "    return sentence\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from MultiCheXNet.data_loader.indiana_dataloader import get_train_validation_generator\n",
    "from tensorflow.keras.applications.densenet import preprocess_input\n",
    "max_vocab_size=10000\n",
    "max_len=100\n",
    "\n",
    "csv_path1  =\"/kaggle/input/chest-xrays-indiana-university/indiana_reports.csv\"\n",
    "csv_path2  =\"/kaggle/input/chest-xrays-indiana-university/indiana_projections.csv\"\n",
    "img_path   =\"/kaggle/input/chest-xrays-indiana-university/images/images_normalized/\"\n",
    "batch_sz = 8\n",
    "validation_split = 0.2\n",
    "\n",
    "train_dataloader, val_dataloader, vocab_size, tok = get_train_validation_generator(csv_path1,csv_path2,img_path, max_vocab_size,max_len, normalize= True,hist_eq=True, augmentation=True, batch_size=batch_sz, validation_split=validation_split,shuffle_GT_sentences=True , over_sample=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GTs = []\n",
    "preds = []\n",
    "for index,(X,Y) in enumerate(val_dataloader):\n",
    "    print(index)\n",
    "    for img,y in zip(X[0],Y):\n",
    "        GT = tokens_to_text(list(y),tok)\n",
    "        pred = get_sentence_preds(img,tok, initial_state_inference_model, inference_model,None)\n",
    "        \n",
    "        GTs.append(GT)\n",
    "        preds.append(pred)\n",
    "        #print(GT)\n",
    "        #print(\"==================================================================\")\n",
    "        #print(pred)\n",
    "\n",
    "\n",
    "        #print(\"\")\n",
    "    if index ==80:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nltk.translate.bleu_score import corpus_bleu\n",
    "def calculate_bleu_evaluation(GT_sentences, predicted_sentences):\n",
    "    BLEU_1 = corpus_bleu(GT_sentences, predicted_sentences, weights=(1.0, 0, 0, 0))\n",
    "    BLEU_2 = corpus_bleu(GT_sentences, predicted_sentences, weights=(0.5, 0.5, 0, 0))\n",
    "    BLEU_3 = corpus_bleu(GT_sentences, predicted_sentences, weights=(0.3, 0.3, 0.3, 0))\n",
    "    BLEU_4 = corpus_bleu(GT_sentences, predicted_sentences, weights=(0.25, 0.25, 0.25, 0.25))\n",
    "    \n",
    "    return BLEU_1,BLEU_2,BLEU_3,BLEU_4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "calculate_bleu_evaluation(GTs,preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}