{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.layers.tcn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TCN\n",
    "> Temporal Convolutional Network.\n",
    "\n",
    "Temporal Convolutional Network (TCN) is a popular architecture. Studies have shown that TCN outperforms canonical recurrent networks such as LSTMs across a diverse range of tasks and datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *\n",
    "from fastcore.nb_imports import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn.utils import weight_norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class Chomp1d(nn.Module):\n",
    "    def __init__(self, chomp_size):\n",
    "        super(Chomp1d, self).__init__()\n",
    "        self.chomp_size = chomp_size\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x[:, :, :-self.chomp_size].contiguous()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class TemporalBlock(nn.Module):\n",
    "    def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, padding, dropout=0.2):\n",
    "        super(TemporalBlock, self).__init__()\n",
    "        self.conv1 = weight_norm(nn.Conv1d(n_inputs, n_outputs, kernel_size,\n",
    "                                           stride=stride, padding=padding, dilation=dilation))\n",
    "        self.chomp1 = Chomp1d(padding)\n",
    "        self.relu1 = nn.ReLU()\n",
    "        self.dropout1 = nn.Dropout(dropout)\n",
    "\n",
    "        self.conv2 = weight_norm(nn.Conv1d(n_outputs, n_outputs, kernel_size,\n",
    "                                           stride=stride, padding=padding, dilation=dilation))\n",
    "        self.chomp2 = Chomp1d(padding)\n",
    "        self.relu2 = nn.ReLU()\n",
    "        self.dropout2 = nn.Dropout(dropout)\n",
    "\n",
    "        self.net = nn.Sequential(self.conv1, self.chomp1, self.relu1, self.dropout1,\n",
    "                                 self.conv2, self.chomp2, self.relu2, self.dropout2)\n",
    "        self.downsample = nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None\n",
    "        self.relu = nn.ReLU()\n",
    "        self.init_weights()\n",
    "\n",
    "    def init_weights(self):\n",
    "        self.conv1.weight.data.normal_(0, 0.01)\n",
    "        self.conv2.weight.data.normal_(0, 0.01)\n",
    "        if self.downsample is not None:\n",
    "            self.downsample.weight.data.normal_(0, 0.01)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.net(x)\n",
    "        res = x if self.downsample is None else self.downsample(x)\n",
    "        return self.relu(out + res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TemporalConvNet(nn.Module):\n",
    "    def __init__(self, num_inputs, num_channels, kernel_size=2, dropout=0.2):\n",
    "        super(TemporalConvNet, self).__init__()\n",
    "        layers = []\n",
    "        num_levels = len(num_channels)\n",
    "        for i in range(num_levels):\n",
    "            dilation_size = 2 ** i\n",
    "            in_channels = num_inputs if i == 0 else num_channels[i-1]\n",
    "            out_channels = num_channels[i]\n",
    "            layers += [TemporalBlock(in_channels, out_channels, kernel_size, stride=1, dilation=dilation_size,\n",
    "                                     padding=(kernel_size-1) * dilation_size, dropout=dropout)]\n",
    "\n",
    "        self.network = nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.network(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **References:-**\n",
    "- Bai, S., Kolter, J. Z., & Koltun, V. (2018). An empirical evaluation of generic convolutional and recurrent networks for sequence modeling. arXiv preprint arXiv:1803.01271.\n",
    "- Official TCN PyTorch implementation: https://github.com/locuslab/TCN\n",
    "- https://arxiv.org/pdf/2112.11023v1.pdf\n",
    "- https://github.com/chenjie04/MPM/blob/master/tcn.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-31 08:12:27\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.144+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "PIL       : 7.1.2\n",
      "IPython   : 5.5.0\n",
      "torch     : 1.10.0+cu111\n",
      "matplotlib: 3.2.2\n",
      "numpy     : 1.19.5\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
