{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.components.tcn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Temporal convolutional network\n",
    "> TCN implementation"
   ]
  },
  {
   "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\n",
    "\n",
    "from nixtlats.models.components.common import Chomp1d\n",
    "from nixtlats.models.components.common import CausalConv1d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "# https://github.com/locuslab/TCN\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 _TemporalBlock2(nn.Module):\n",
    "    def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, padding, dropout=0.2):\n",
    "        super(_TemporalBlock2, self).__init__()\n",
    "        self.causalconv1 = CausalConv1d(in_channels=n_inputs, out_channels=n_outputs, \n",
    "                                        kernel_size=kernel_size, stride=stride, padding=padding, \n",
    "                                        dilation=dilation, activation='ReLU', with_weight_norm=True)\n",
    "\n",
    "        self.causalconv2 = CausalConv1d(in_channels=n_outputs, out_channels=n_outputs, \n",
    "                                        kernel_size=kernel_size, stride=stride, padding=padding, \n",
    "                                        dilation=dilation, activation='ReLU', with_weight_norm=True)\n",
    "\n",
    "        self.net = nn.Sequential(self.causalconv1, nn.Dropout(dropout),\n",
    "                                 self.causalconv2, nn.Dropout(dropout))\n",
    "        \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.causalconv1.conv.weight.data.normal_(0, 0.01)\n",
    "        self.causalconv2.conv.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",
    "            layers += [_TemporalBlock2(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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from torchsummary import summary\n",
    "\n",
    "# model = _TemporalConvNet(num_inputs=7*24, num_channels=[2,2,2], kernel_size=2, dropout=0.2)\n",
    "# summary(model, (1, 7*24))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
