Signal2Signal.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Signal2Signal with RNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import the required packages :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook we train a Recurrent Neural Network to predict a signal to signal correspondance. More precisely, we want to reproduce the following sequence:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$(f_1(t_i),f_2(t_i),f_3(t_i))\\to(f_4(t_i) = f_1(t_i)-f_2(t_i)-f_3(t_i))$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where $f_i$ are given functions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The functions are given by:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f1(t):\n",
    "   \"f1 function\"\n",
    "   ret = torch.sin(t)\n",
    "   return ret\n",
    "def f2(t):\n",
    "   \"f2 function\"\n",
    "   ret = torch.cos(t)\n",
    "   return ret\n",
    "def f3(t):\n",
    "   \"f3 function\"\n",
    "   ret = t\n",
    "   return ret\n",
    "def f4(t):\n",
    "   \"f4 function\"\n",
    "   ret = f1(t) - f2(t) - f3(t) \n",
    "   return ret"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "They look like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dt = 0.1\n",
    "t = torch.arange(-1.0, 1.0, dt)\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(t.numpy(),f1(t).numpy(),\"-\", label=\"f1\")\n",
    "ax.plot(t.numpy(),f2(t).numpy(), \"-\", label=\"f2\")\n",
    "ax.plot(t.numpy(),f3(t).numpy(), \"-\", label=\"f3\")\n",
    "ax.plot(t.numpy(),f4(t).numpy(), \"-\", label=\"f4\")\n",
    "ax.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will cut the data in sequences of size seqlen data points. The input and target sequences are:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\\left(\n",
    "\\begin{array}{cccc}\n",
    " f_1(t_{i}) & f_2(t_{i}) & f_3(t_{i})\\\\\n",
    " \\vdots  & \\vdots  & \\vdots  \\\\\n",
    " f_1(t_{i+seqlen}) & f_2(t_{i+seqlen}) & f_3(t_{i+seqlen})\n",
    "\\end{array}\n",
    "\\right)\n",
    "\\to\n",
    "\\left(\n",
    "\\begin{array}{cccc}\n",
    "f_4(t_{i}) \\\\\n",
    "\\vdots  \\\\\n",
    "f_4(t_{i+seqlen})\n",
    "\\end{array}\n",
    "\\right)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create the Dataset object to manage the data. The method ``__getitem__`` returns a training example of the form:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\\left(\n",
    "\\begin{array}{cccc}\n",
    " x_1^1 & x_1^2 & \\ldots & x_1^{inputSize_x} \\\\\n",
    " \\vdots  & \\vdots  & \\ldots  & \\vdots  \\\\\n",
    " x_{seqlen}^1 & x_{seqlen}^2 & \\ldots & x_{seqlen}^{inputSize_x} \\\\\n",
    "\\end{array}\n",
    "\\right)\n",
    "\\to\n",
    "\\left(\n",
    "\\begin{array}{cccc}\n",
    " y_1^1 & y_1^2 & \\ldots & y_1^{inputSize_y} \\\\\n",
    " \\vdots  & \\vdots  & \\ldots  & \\vdots  \\\\\n",
    " y_{seqlen}^1 & y_{seqlen}^2 & \\ldots & y_{seqlen}^{inputSize_y} \\\\\n",
    "\\end{array}\n",
    "\\right)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sig2SigDataset(Dataset):\n",
    "    \"\"\" Sig2Sig dataset.\"\"\"\n",
    "    \n",
    "    def __init__(self, h_size, seq_len):\n",
    "        \n",
    "        # Generate the data\n",
    "        tt = torch.arange(-1.0, 1.0, h_size)  \n",
    "        lentt = tt.shape[0]\n",
    "        \n",
    "        # Check if h_size and seq_len are compatible\n",
    "        if lentt - seq_len < 0:\n",
    "            raise TypeError(f'h_size = {h_size} is too large for the seq_len = {seq_len}')     \n",
    "        \n",
    "        T1 = f1(tt).view([1,lentt]).t()\n",
    "        T2 = f2(tt).view([1,lentt]).t()\n",
    "        T3 = f3(tt).view([1,lentt]).t()\n",
    "        dataX = torch.cat((T1, T2, T3), 1)    \n",
    "        T4 = f4(tt).view([1,lentt]).t()\n",
    "        dataY = T4 \n",
    "        \n",
    "        self.X = torch.zeros(lentt - seq_len, seq_len, 3)\n",
    "        self.Y = torch.zeros(lentt - seq_len, seq_len, 1)\n",
    "        \n",
    "        self.len = self.X.shape[0]\n",
    "        \n",
    "        for i in range(self.len):\n",
    "            self.X[i,:,:] = dataX[i:i + seq_len,:]\n",
    "            self.Y[i,:,:] = dataY[i:i + seq_len,:]\n",
    "                \n",
    "    def __getitem__(self, index):\n",
    "        return self.X[index,:,:], self.Y[index,:,:]\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.len"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Design the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Using `torch.nn.RNN` or `torch.nn.LSTM`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We creeate a RNN with a readout layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNNModel(nn.Module):\n",
    "    \n",
    "    def __init__(self, input_dim, hidden_dim, layer_dim, output_dim, rnn_type = 'RNN'):\n",
    "        super().__init__()\n",
    "        \n",
    "        # Number of hidden dimensions\n",
    "        self.hidden_dim = hidden_dim\n",
    "        \n",
    "        # Number of hidden layers\n",
    "        self.layer_dim = layer_dim\n",
    "        \n",
    "        # RNN\n",
    "        if rnn_type == 'RNN':\n",
    "            self.rnn = nn.RNN(input_dim, hidden_dim, layer_dim, batch_first = True, nonlinearity = 'tanh')\n",
    "        elif rnn_type == 'LSTM':    \n",
    "            self.rnn = nn.LSTM(input_dim, hidden_dim, layer_dim, batch_first = True)\n",
    "        \n",
    "        # Readout layer\n",
    "        self.fc = nn.Linear(hidden_dim, output_dim)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # Initialize hidden state with zeros by default         \n",
    "        # One time step\n",
    "        out, hn = self.rnn(x)\n",
    "        out = self.fc(out) \n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Using `torch.nn.RNNCell` or `torch.nn.LSTMCell`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We creeate a RNN with a cell layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNNModelWithCell(nn.Module):\n",
    "    \n",
    "    def __init__(self, input_dim, hidden_dim, output_dim, rnn_type = 'RNN'):\n",
    "        super().__init__()\n",
    "        \n",
    "        # Number of hidden dimensions\n",
    "        self.hidden_dim = hidden_dim\n",
    "        \n",
    "        # Cell type\n",
    "        self.rnn_type = rnn_type\n",
    "        \n",
    "        # RNN\n",
    "        if self.rnn_type == 'RNN':\n",
    "            self.rnn = nn.RNNCell(input_dim, hidden_dim, nonlinearity = 'tanh')\n",
    "        elif self.rnn_type == 'LSTM':    \n",
    "            self.rnn = nn.LSTMCell(input_dim, hidden_dim)\n",
    "        \n",
    "        # Readout layer\n",
    "        self.fc = nn.Linear(hidden_dim, output_dim)\n",
    "        \n",
    "    def forward(self, x):    \n",
    "        output = []\n",
    "        # Initialize hidden state with zeros\n",
    "        # hn.shape must be torch.Size([len_seq, hidden_dim])\n",
    "        hn = torch.zeros(x.size(1), self.hidden_dim)\n",
    "        if self.rnn_type == 'LSTM':\n",
    "            cx = torch.zeros(x.size(1), self.hidden_dim)\n",
    "        for i in range(x.size(0)):\n",
    "            # One time step\n",
    "            if self.rnn_type == 'RNN':\n",
    "                hn = self.rnn(x[i], hn)\n",
    "            elif self.rnn_type == 'LSTM':    \n",
    "                hn, cx = self.rnn(x[i], (hn, cx))\n",
    "            out = self.fc(hn)\n",
    "            output.append(out)\n",
    "        return torch.stack(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialization Algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a initialization weight function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialitation\n",
    "def init_weights(self):\n",
    "    for idx, m in enumerate(self.modules()):         \n",
    "        if idx > 0:\n",
    "            print('{} -> {}'.format(idx,m))    \n",
    "            if type(m) in [nn.RNN, nn.GRU, nn.LSTM, nn.RNNCell, nn.GRUCell, nn.LSTMCell]:\n",
    "                for name, param in m.named_parameters():\n",
    "                    print(f'Initialization of {name}', end=\"\", flush=True)\n",
    "                    if 'weight_ih' in name:   \n",
    "                        torch.nn.init.xavier_uniform_(param.data)\n",
    "                        print('...done')\n",
    "                    elif 'weight_hh' in name:\n",
    "                        torch.nn.init.orthogonal_(param.data)\n",
    "                        print('...done')\n",
    "                    elif 'bias' in name:\n",
    "                        param.data.fill_(0)\n",
    "                        print('...done')\n",
    "            if type(m) in [nn.Linear]:\n",
    "                for name, param in m.named_parameters():\n",
    "                    print(f'Initialization of {name}', end=\"\", flush=True)\n",
    "                    if 'weight' in name:   \n",
    "                        torch.nn.init.xavier_uniform_(param.data)\n",
    "                        print('...done')\n",
    "                    elif 'bias' in name:\n",
    "                        param.data.fill_(0)\n",
    "                        print('...done')           "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup the training process: Model + Loss + Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data parameters\n",
    "h_size = 0.001\n",
    "seq_len = 100\n",
    "\n",
    "# Global parameters\n",
    "input_dim = 3      # input dimension\n",
    "hidden_dim = 100   # hidden layer dimension\n",
    "layer_dim = 1      # number of hidden layers\n",
    "output_dim = 1     # output dimension\n",
    "rnn_type = 'RNN'   # type of cell\n",
    "cell_model = True  # Wether we use nn.RNN or nn.RNNCell\n",
    "\n",
    "# Dataset\n",
    "dataset = Sig2SigDataset(h_size, seq_len)\n",
    "print('Length of dataset: {}\\n'.format(dataset.__len__()))\n",
    "# DataLoader\n",
    "train_loader = DataLoader(dataset = dataset,\n",
    "                          batch_size = 32,\n",
    "                          shuffle = True,\n",
    "                          num_workers = 1)\n",
    "\n",
    "# Model\n",
    "if not cell_model:\n",
    "    model = RNNModel(input_dim, hidden_dim, layer_dim, output_dim, rnn_type)\n",
    "else:\n",
    "    model = RNNModelWithCell(input_dim, hidden_dim, output_dim, rnn_type)\n",
    "    \n",
    "# Initialization                  \n",
    "model.apply(init_weights)\n",
    "\n",
    "# Criterion\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "# Optimizer\n",
    "learning_rate = 0.01\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate)\n",
    "\n",
    "# Scheduler\n",
    "gamma = 0.5 # Decay LR by a factor of gamma every step_size epochs\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size = 7, gamma = gamma)\n",
    "#scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Create the training loop:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 500\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    \n",
    "    # Adjust the learning rate\n",
    "    scheduler.step()\n",
    "    \n",
    "    for i, data in enumerate(train_loader, 0):\n",
    " \n",
    "        # Get the inputs\n",
    "        X, y = data\n",
    "\n",
    "        # Train step\n",
    "        model.train()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        y_ = model(X)\n",
    "        loss = criterion(y_, y)\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if epoch % 100 == 0:\n",
    "            if i % 20 == 0:        \n",
    "                # Evaluation\n",
    "                model.eval()\n",
    "                with torch.no_grad():\n",
    "                    y_ = model(X)\n",
    "\n",
    "                # Compute and print loss\n",
    "                loss = criterion(y_, y)\n",
    "                \n",
    "                print(f'Epoch : {epoch}, Iteration : {i}, Loss : {loss.item()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Post-processing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Investigate the result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualization\n",
    "with torch.no_grad():\n",
    "    \n",
    "    h = 0.001;\n",
    "    tt = torch.arange(-1.0, 1.0, h)  \n",
    "    lentt = tt.shape[0]   \n",
    "    T1 = f1(tt).view([1,lentt]).t()\n",
    "    T2 = f2(tt).view([1,lentt]).t()\n",
    "    T3 = f3(tt).view([1,lentt]).t()\n",
    "    dataX = torch.cat((T1, T2, T3), 1)\n",
    "    dataX = dataX.view(-1, dataX.shape[0], dataX.shape[1])\n",
    "    T4 = f4(tt).view([1,lentt]).t()\n",
    "    data = T4 \n",
    "    \n",
    "    # Model in evaluation state\n",
    "    model.eval()      \n",
    "       \n",
    "    # Evalaute the model\n",
    "    y_ = model(dataX)\n",
    "    \n",
    "    fig, ax = plt.subplots()\n",
    "    ax.plot(data.numpy(),\".\", label=\"data\")\n",
    "    ax.plot(y_.view([lentt,1]).numpy(), \".\", label=\"pred\")\n",
    "    ax.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6.4 64-bit ('3.6.4')",
   "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.6.4"
  },
  "vscode": {
   "interpreter": {
    "hash": "185474264c9975e99e952361575b58ce6ef8c0725efde3d0ce11f857d91eb0e2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}