{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.components.drnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev import *\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dilated Recurrent Neural Networks\n",
    "\n",
    "> DRNN pytorch implementation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dilated Recurrent Neural Networks. https://papers.nips.cc/paper/6613-dilated-recurrent-neural-networks.pdf\n",
    "implementation from https://github.com/zalandoresearch/pytorch-dilated-rnn\n",
    "Residual LSTM: Design of a Deep Recurrent Architecture for Distant Speech Recognition. https://arxiv.org/abs/1701.03360\n",
    "A Dual-Stage Attention-Based recurrent neural network for time series prediction. https://arxiv.org/abs/1704.02971\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.autograd as autograd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class LSTMCell(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, dropout=0.):\n",
    "        super(LSTMCell, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.weight_ih = nn.Parameter(torch.randn(4 * hidden_size, input_size))\n",
    "        self.weight_hh = nn.Parameter(torch.randn(4 * hidden_size, hidden_size))\n",
    "        self.bias_ih = nn.Parameter(torch.randn(4 * hidden_size))\n",
    "        self.bias_hh = nn.Parameter(torch.randn(4 * hidden_size))\n",
    "        self.dropout = dropout\n",
    "\n",
    "    def forward(self, inputs, hidden):\n",
    "        hx, cx = hidden[0].squeeze(0), hidden[1].squeeze(0)\n",
    "        gates = (torch.matmul(inputs, self.weight_ih.t()) + self.bias_ih +\n",
    "                         torch.matmul(hx, self.weight_hh.t()) + self.bias_hh)\n",
    "        ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)\n",
    "\n",
    "        ingate = torch.sigmoid(ingate)\n",
    "        forgetgate = torch.sigmoid(forgetgate)\n",
    "        cellgate = torch.tanh(cellgate)\n",
    "        outgate = torch.sigmoid(outgate)\n",
    "\n",
    "        cy = (forgetgate * cx) + (ingate * cellgate)\n",
    "        hy = outgate * torch.tanh(cy)\n",
    "\n",
    "        return hy, (hy, cy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ResLSTMCell(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, dropout=0.):\n",
    "        super(ResLSTMCell, self).__init__()\n",
    "        self.register_buffer('input_size', torch.Tensor([input_size]))\n",
    "        self.register_buffer('hidden_size', torch.Tensor([hidden_size]))\n",
    "        self.weight_ii = nn.Parameter(torch.randn(3 * hidden_size, input_size))\n",
    "        self.weight_ic = nn.Parameter(torch.randn(3 * hidden_size, hidden_size))\n",
    "        self.weight_ih = nn.Parameter(torch.randn(3 * hidden_size, hidden_size))\n",
    "        self.bias_ii = nn.Parameter(torch.randn(3 * hidden_size))\n",
    "        self.bias_ic = nn.Parameter(torch.randn(3 * hidden_size))\n",
    "        self.bias_ih = nn.Parameter(torch.randn(3 * hidden_size))\n",
    "        self.weight_hh = nn.Parameter(torch.randn(1 * hidden_size, hidden_size))\n",
    "        self.bias_hh = nn.Parameter(torch.randn(1 * hidden_size))\n",
    "        self.weight_ir = nn.Parameter(torch.randn(hidden_size, input_size))\n",
    "        self.dropout = dropout\n",
    "\n",
    "    def forward(self, inputs, hidden):\n",
    "        hx, cx = hidden[0].squeeze(0), hidden[1].squeeze(0)\n",
    "\n",
    "        ifo_gates = (torch.matmul(inputs, self.weight_ii.t()) + self.bias_ii +\n",
    "                                  torch.matmul(hx, self.weight_ih.t()) + self.bias_ih +\n",
    "                                  torch.matmul(cx, self.weight_ic.t()) + self.bias_ic)\n",
    "        ingate, forgetgate, outgate = ifo_gates.chunk(3, 1)\n",
    "\n",
    "        cellgate = torch.matmul(hx, self.weight_hh.t()) + self.bias_hh\n",
    "\n",
    "        ingate = torch.sigmoid(ingate)\n",
    "        forgetgate = torch.sigmoid(forgetgate)\n",
    "        cellgate = torch.tanh(cellgate)\n",
    "        outgate = torch.sigmoid(outgate)\n",
    "\n",
    "        cy = (forgetgate * cx) + (ingate * cellgate)\n",
    "        ry = torch.tanh(cy)\n",
    "\n",
    "        if self.input_size == self.hidden_size:\n",
    "            hy = outgate * (ry + inputs)\n",
    "        else:\n",
    "            hy = outgate * (ry + torch.matmul(inputs, self.weight_ir.t()))\n",
    "        return hy, (hy, cy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ResLSTMLayer(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, dropout=0.):\n",
    "        super(ResLSTMLayer, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.cell = ResLSTMCell(input_size, hidden_size, dropout=0.)\n",
    "\n",
    "    def forward(self, inputs, hidden):\n",
    "        inputs = inputs.unbind(0)\n",
    "        outputs = []\n",
    "        for i in range(len(inputs)):\n",
    "                out, hidden = self.cell(inputs[i], hidden)\n",
    "                outputs += [out]\n",
    "        outputs = torch.stack(outputs)\n",
    "        return outputs, hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class AttentiveLSTMLayer(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, dropout=0.0):\n",
    "        super(AttentiveLSTMLayer, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        attention_hsize = hidden_size\n",
    "        self.attention_hsize = attention_hsize\n",
    "\n",
    "        self.cell = LSTMCell(input_size, hidden_size)\n",
    "        self.attn_layer = nn.Sequential(nn.Linear(2 * hidden_size + input_size, attention_hsize),\n",
    "                                        nn.Tanh(),\n",
    "                                        nn.Linear(attention_hsize, 1))\n",
    "        self.softmax = nn.Softmax(dim=0)\n",
    "        self.dropout = dropout\n",
    "\n",
    "    def forward(self, inputs, hidden):\n",
    "        inputs = inputs.unbind(0)\n",
    "        outputs = []\n",
    "\n",
    "        for t in range(len(inputs)):\n",
    "            # attention on windows\n",
    "            hx, cx = (tensor.squeeze(0) for tensor in hidden)\n",
    "            hx_rep = hx.repeat(len(inputs), 1, 1)\n",
    "            cx_rep = cx.repeat(len(inputs), 1, 1)\n",
    "            x = torch.cat((inputs, hx_rep, cx_rep), dim=-1)\n",
    "            l = self.attn_layer(x)\n",
    "            beta = self.softmax(l)\n",
    "            context = torch.bmm(beta.permute(1, 2, 0),\n",
    "                                inputs.permute(1, 0, 2)).squeeze(1)\n",
    "            out, hidden = self.cell(context, hidden)\n",
    "            outputs += [out]\n",
    "        outputs = torch.stack(outputs)\n",
    "        return outputs, hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class DRNN(nn.Module):\n",
    "\n",
    "    def __init__(self, n_input, n_hidden, n_layers, dilations, dropout=0, cell_type='GRU', batch_first=False):\n",
    "        super(DRNN, self).__init__()\n",
    "\n",
    "        self.dilations = dilations\n",
    "        self.cell_type = cell_type\n",
    "        self.batch_first = batch_first\n",
    "\n",
    "        layers = []\n",
    "        if self.cell_type == \"GRU\":\n",
    "            cell = nn.GRU\n",
    "        elif self.cell_type == \"RNN\":\n",
    "            cell = nn.RNN\n",
    "        elif self.cell_type == \"LSTM\":\n",
    "            cell = nn.LSTM\n",
    "        elif self.cell_type == \"ResLSTM\":\n",
    "            cell = ResLSTMLayer\n",
    "        elif self.cell_type == \"AttentiveLSTM\":\n",
    "            cell = AttentiveLSTMLayer\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "\n",
    "        for i in range(n_layers):\n",
    "            if i == 0:\n",
    "                c = cell(n_input, n_hidden, dropout=dropout)\n",
    "            else:\n",
    "                c = cell(n_hidden, n_hidden, dropout=dropout)\n",
    "            layers.append(c)\n",
    "        self.cells = nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, inputs, hidden=None):\n",
    "        if self.batch_first:\n",
    "            inputs = inputs.transpose(0, 1)\n",
    "        outputs = []\n",
    "        for i, (cell, dilation) in enumerate(zip(self.cells, self.dilations)):\n",
    "            if hidden is None:\n",
    "                inputs, _ = self.drnn_layer(cell, inputs, dilation)\n",
    "            else:\n",
    "                inputs, hidden[i] = self.drnn_layer(cell, inputs, dilation, hidden[i])\n",
    "\n",
    "            outputs.append(inputs[-dilation:])\n",
    "\n",
    "        if self.batch_first:\n",
    "            inputs = inputs.transpose(0, 1)\n",
    "        return inputs, outputs\n",
    "\n",
    "    def drnn_layer(self, cell, inputs, rate, hidden=None):\n",
    "        n_steps = len(inputs)\n",
    "        batch_size = inputs[0].size(0)\n",
    "        hidden_size = cell.hidden_size\n",
    "\n",
    "        inputs, dilated_steps = self._pad_inputs(inputs, n_steps, rate)\n",
    "        dilated_inputs = self._prepare_inputs(inputs, rate)\n",
    "\n",
    "        if hidden is None:\n",
    "            dilated_outputs, hidden = self._apply_cell(dilated_inputs, cell, batch_size, rate, hidden_size)\n",
    "        else:\n",
    "            hidden = self._prepare_inputs(hidden, rate)\n",
    "            dilated_outputs, hidden = self._apply_cell(dilated_inputs, cell, batch_size, rate, hidden_size,\n",
    "                                                       hidden=hidden)\n",
    "\n",
    "        splitted_outputs = self._split_outputs(dilated_outputs, rate)\n",
    "        outputs = self._unpad_outputs(splitted_outputs, n_steps)\n",
    "\n",
    "        return outputs, hidden\n",
    "\n",
    "    def _apply_cell(self, dilated_inputs, cell, batch_size, rate, hidden_size, hidden=None):\n",
    "        if hidden is None:\n",
    "            hidden = torch.zeros(batch_size * rate, hidden_size,\n",
    "                                 dtype=dilated_inputs.dtype,\n",
    "                                 device=dilated_inputs.device)\n",
    "            hidden = hidden.unsqueeze(0)\n",
    "            \n",
    "            if self.cell_type in ['LSTM', 'ResLSTM', 'AttentiveLSTM']:\n",
    "                hidden = (hidden, hidden)\n",
    "                \n",
    "        dilated_outputs, hidden = cell(dilated_inputs, hidden) # compatibility hack\n",
    "\n",
    "        return dilated_outputs, hidden\n",
    "\n",
    "    def _unpad_outputs(self, splitted_outputs, n_steps):\n",
    "        return splitted_outputs[:n_steps]\n",
    "\n",
    "    def _split_outputs(self, dilated_outputs, rate):\n",
    "        batchsize = dilated_outputs.size(1) // rate\n",
    "\n",
    "        blocks = [dilated_outputs[:, i * batchsize: (i + 1) * batchsize, :] for i in range(rate)]\n",
    "\n",
    "        interleaved = torch.stack((blocks)).transpose(1, 0).contiguous()\n",
    "        interleaved = interleaved.view(dilated_outputs.size(0) * rate,\n",
    "                                       batchsize,\n",
    "                                       dilated_outputs.size(2))\n",
    "        return interleaved\n",
    "\n",
    "    def _pad_inputs(self, inputs, n_steps, rate):\n",
    "        iseven = (n_steps % rate) == 0\n",
    "\n",
    "        if not iseven:\n",
    "            dilated_steps = n_steps // rate + 1\n",
    "\n",
    "            zeros_ = torch.zeros(dilated_steps * rate - inputs.size(0),\n",
    "                                 inputs.size(1),\n",
    "                                 inputs.size(2), \n",
    "                                 dtype=inputs.dtype,\n",
    "                                 device=inputs.device)\n",
    "            inputs = torch.cat((inputs, zeros_))\n",
    "        else:\n",
    "            dilated_steps = n_steps // rate\n",
    "\n",
    "        return inputs, dilated_steps\n",
    "\n",
    "    def _prepare_inputs(self, inputs, rate):\n",
    "        dilated_inputs = torch.cat([inputs[j::rate, :, :] for j in range(rate)], 1)\n",
    "        return dilated_inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_time = 24\n",
    "batch_size = 3\n",
    "n_feats = 10\n",
    "n_output_feats = 20\n",
    "num_layers = 4\n",
    "dilations = [2, 4, 5, 9]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test forward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for cell_type in ['GRU', 'RNN', 'LSTM', 'ResLSTM']:\n",
    "    model = DRNN(n_input=n_feats, n_hidden=n_output_feats, n_layers=num_layers, \n",
    "             dilations=dilations, cell_type=cell_type)\n",
    "    x = torch.rand(n_time, batch_size, n_feats)\n",
    "    output, hidden = model(x)\n",
    "    test_eq(output.shape, (n_time, batch_size, n_output_feats))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
