RainPredNet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8EQyKmQUYADk"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.optim as optim\n",
    "\n",
    "import glob\n",
    "from PIL import Image\n",
    "import time\n",
    "import os\n",
    "import h5py\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jYM_lu9tYADs"
   },
   "outputs": [],
   "source": [
    "# Define some constants\n",
    "KERNEL_SIZE = 3\n",
    "PADDING = KERNEL_SIZE//2\n",
    "BATCH_SIZE = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ID5fszzOYADn"
   },
   "outputs": [],
   "source": [
    "colors = [\n",
    "   (255, 255, 255),\n",
    "   (236, 254, 252),\n",
    "   (196, 254, 252),\n",
    "   (156, 234, 252),\n",
    "   (156, 218, 252),\n",
    "   (180, 198, 252),\n",
    "   (180, 254, 156),\n",
    "   (180, 234, 156),\n",
    "   (164, 218, 156),\n",
    "   (252, 254, 156),\n",
    "   (252, 234, 156),\n",
    "   (252, 218, 156),\n",
    "   (252, 190, 196),\n",
    "   (252, 158, 156),\n",
    "   (228, 158, 164)\n",
    "]\n",
    "\n",
    "amount = [\n",
    "    0,\n",
    "    1,\n",
    "    2,\n",
    "    4,\n",
    "    8,\n",
    "    12,\n",
    "    16,\n",
    "    24,\n",
    "    32,\n",
    "    40,\n",
    "    48,\n",
    "    56,\n",
    "    64,\n",
    "    80,\n",
    "    100,\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bTMSbzaEVwOh"
   },
   "source": [
    "# Prepare Dataset and Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "63_xJv4MdE1l"
   },
   "outputs": [],
   "source": [
    "rainy_count = pd.read_csv('rain_count.csv', header=None)\n",
    "rainy_count.columns = ['date', 'rain']\n",
    "rainy_count['date'] = pd.to_datetime(rainy_count['date'])\n",
    "rainy_count.set_index('date', inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0ldNTIgDdYyS"
   },
   "outputs": [],
   "source": [
    "rain_threshold = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "u5vuikdmdnW9"
   },
   "outputs": [],
   "source": [
    "rainy_hours = rainy_count[rainy_count.rain > rain_threshold].index[1:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dmW7fEAvdtcE"
   },
   "outputs": [],
   "source": [
    "train_times = rainy_hours[rainy_hours.year < 2017]\n",
    "valid_times = rainy_hours[rainy_hours.year == 2017]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "4xi4KvLzo0Q0"
   },
   "outputs": [],
   "source": [
    "ts = pd.date_range('2013', '2016', freq='5min')\n",
    "train_idx = np.empty(len(train_times), dtype=int)\n",
    "for i,t in enumerate(train_times):\n",
    "    idx = ts.get_loc(t)\n",
    "    train_idx[i] = idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_w0GsqTInzaG"
   },
   "outputs": [],
   "source": [
    "ts = pd.date_range('2016', '2017', freq='5min')\n",
    "valid_idx = np.empty(len(valid_times), dtype=int)\n",
    "for i,t in enumerate(valid_times):\n",
    "    idx = ts.get_loc(t)\n",
    "    valid_idx[i] = idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RLOzraEYpFbR"
   },
   "outputs": [],
   "source": [
    "t = train_times[0]-pd.Timedelta('1h')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "LJAV0MKi--Xa",
    "outputId": "bb5b31e5-66a7-49e8-9007-12547a6dcf8a"
   },
   "outputs": [],
   "source": [
    "filename = \"{:4d}{:02d}/{:4d}{:02d}{:02d}{:02d}{:02d}.npz\".format(t.year,t.month,t.year,t.month,t.day,t.hour,t.minute)\n",
    "print(filename)\n",
    "a = np.load(filename)['arr_0']\n",
    "print(a.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nw0yh-_MBluc"
   },
   "outputs": [],
   "source": [
    "b = np.zeros((24,1,a.shape[0],a.shape[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "6LICMepeCdQm",
    "outputId": "14bc0fc1-b299-4320-ea5f-26d09deed518"
   },
   "outputs": [],
   "source": [
    "b.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lreXOYleCxNJ"
   },
   "outputs": [],
   "source": [
    "steps = 24"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 425
    },
    "colab_type": "code",
    "id": "B-4otxkvCip0",
    "outputId": "497ee4fb-7ad6-4bf4-d162-3bf337a4fe06"
   },
   "outputs": [],
   "source": [
    "t1 = train_times[100]\n",
    "t0 = t1 - steps//2 * pd.Timedelta('5min')\n",
    "t2 = t1 + steps//2 * pd.Timedelta('5min')\n",
    "#print(t1)\n",
    "for i,t in enumerate(pd.date_range(t0, t2, freq='5min')[:-1]):\n",
    "    filename = \"{:4d}{:02d}/{:4d}{:02d}{:02d}{:02d}{:02d}.npz\".format(t.year,t.month,t.year,t.month,t.day,t.hour,t.minute)\n",
    "    print(i, filename)\n",
    "    img = np.load(filename)['arr_0']\n",
    "    b[i,0,:,:] = img\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 687
    },
    "colab_type": "code",
    "id": "bDfLZUc9DSY5",
    "outputId": "80a6a5a6-b960-4c83-d8be-9a65158863ff"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(16,12))\n",
    "for i in range(12):\n",
    "    plt.subplot(3,4,i+1)\n",
    "    plt.imshow(b[i,0], vmax=7)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "s2WdWegHavKE"
   },
   "outputs": [],
   "source": [
    "class KyotoRadarDataset(Dataset):\n",
    "    \"\"\"Kyoto Radar dataset.\"\"\"\n",
    "\n",
    "    def __init__(self, times, steps):\n",
    "        \"\"\"\n",
    "        \"\"\"\n",
    "        self.times = times\n",
    "        self.steps = steps\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.times) #- self.steps\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        out = np.zeros((self.steps,1,96,96), dtype=np.float32)\n",
    "        t1 = self.times[idx]\n",
    "        t0 = t1 - self.steps//2 * pd.Timedelta('5min')\n",
    "        t2 = t1 + self.steps//2 * pd.Timedelta('5min')\n",
    "        for i,t in enumerate(pd.date_range(t0, t2, freq='5min')[:-1]):\n",
    "            filename = \"{:4d}{:02d}/{:4d}{:02d}{:02d}{:02d}{:02d}.npz\".format(t.year,t.month,t.year,t.month,t.day,t.hour,t.minute)\n",
    "            img = np.load(filename)['arr_0']\n",
    "            out[i,0,:,:] = img\n",
    "        out = out / 15.\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 704
    },
    "colab_type": "code",
    "id": "bohxW9jpdOsB",
    "outputId": "162043cf-0c37-46bf-e5cd-f122ccfc37ba"
   },
   "outputs": [],
   "source": [
    "dataset = KyotoRadarDataset(times=train_times, steps=12)\n",
    "print(train_times[100])\n",
    "plt.figure(figsize=(16,12))\n",
    "s = dataset[100]\n",
    "for i in range(12):\n",
    "    plt.subplot(3,4,i+1)\n",
    "    plt.imshow(s[i,0], vmax=0.4)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MOifbK4c_vdF"
   },
   "outputs": [],
   "source": [
    "radardata_train = KyotoRadarDataset(times=train_times, steps=24)\n",
    "radardata_valid = KyotoRadarDataset(times=valid_times, steps=24)\n",
    "\n",
    "train_loader = DataLoader(dataset = radardata_train, batch_size = BATCH_SIZE, shuffle=False)\n",
    "valid_loader = DataLoader(dataset = radardata_valid, batch_size = BATCH_SIZE, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "xg9MIxwvYAEV",
    "outputId": "f5131712-9f86-49fb-9566-49bd84c64091"
   },
   "outputs": [],
   "source": [
    "#%%time\n",
    "# obtain one batch of training images\n",
    "dataiter = iter(train_loader)\n",
    "images = dataiter.next()\n",
    "images.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 550
    },
    "colab_type": "code",
    "id": "lGUkqLOtYAEZ",
    "outputId": "a51ccaf8-33d5-42ca-ac01-f0b16f8ebd82"
   },
   "outputs": [],
   "source": [
    "images = dataiter.next()\n",
    "plt.figure(figsize=(12,8))\n",
    "bn = 0\n",
    "for i in range(24):\n",
    "    plt.subplot(4,6,i+1)\n",
    "    plt.tight_layout()\n",
    "    plt.imshow(images[bn,i,0], interpolation='none')\n",
    "    plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OrzzgEPkYAD0"
   },
   "outputs": [],
   "source": [
    "class ConvLSTMCell(nn.Module):\n",
    "\n",
    "    def __init__(self, input_dim, hidden_dim, kernel_size, bias):\n",
    "        super(ConvLSTMCell, self).__init__()\n",
    "\n",
    "        self.input_dim  = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        self.kernel_size = kernel_size\n",
    "        self.padding     = kernel_size // 2\n",
    "        self.bias        = bias\n",
    "        \n",
    "        self.conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,\n",
    "                              out_channels=4 * self.hidden_dim,\n",
    "                              kernel_size=self.kernel_size,\n",
    "                              padding=self.padding,\n",
    "                              bias=self.bias)\n",
    "\n",
    "    def forward(self, input_tensor, cur_state):\n",
    "        h_cur, c_cur = cur_state\n",
    "        \n",
    "        combined = torch.cat([input_tensor, h_cur], dim=1)\n",
    "        \n",
    "        combined_conv = self.conv(combined)\n",
    "        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1) \n",
    "        i = torch.sigmoid(cc_i)\n",
    "        f = torch.sigmoid(cc_f)\n",
    "        o = torch.sigmoid(cc_o)\n",
    "        g = torch.tanh(cc_g)\n",
    "\n",
    "        c_next = f * c_cur + i * g\n",
    "        h_next = o * torch.tanh(c_next)\n",
    "        \n",
    "        return h_next, (h_next, c_next)\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.conv.reset_parameters()\n",
    "\n",
    "    def init_hidden(self, batch_size, width, height):\n",
    "        return (torch.zeros(batch_size, self.hidden_dim, width, height).cuda(),\n",
    "                torch.zeros(batch_size, self.hidden_dim, width, height).cuda())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lzVeyldXYAD4"
   },
   "outputs": [],
   "source": [
    "class SatLU(nn.Module):\n",
    "    def __init__(self, lower=0, upper=15/15, inplace=False):\n",
    "        super(SatLU, self).__init__()\n",
    "        self.lower = lower\n",
    "        self.upper = upper\n",
    "        self.inplace = inplace\n",
    "\n",
    "    def forward(self, input):\n",
    "        return F.hardtanh(input, self.lower, self.upper, self.inplace)\n",
    "\n",
    "    def __repr__(self):\n",
    "        inplace_str = ', inplace' if self.inplace else ''\n",
    "        return self.__class__.__name__ + ' ('\\\n",
    "            + 'min_val=' + str(self.lower) \\\n",
    "            + ', max_val=' + str(self.upper) \\\n",
    "            + inplace_str + ')'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LH3j9KBAYAD6"
   },
   "outputs": [],
   "source": [
    "class PredNet(nn.Module):\n",
    "    def __init__(self, R_channels, A_channels):\n",
    "        super(PredNet, self).__init__()\n",
    "        self.r_channels = R_channels + (0, )  # for convenience\n",
    "        self.a_channels = A_channels\n",
    "        self.n_layers = len(R_channels)\n",
    "\n",
    "        for i in range(self.n_layers):\n",
    "            cell = ConvLSTMCell(2 * self.a_channels[i] + self.r_channels[i+1], \n",
    "                                self.r_channels[i], KERNEL_SIZE, True)\n",
    "\n",
    "            setattr(self, 'cell{}'.format(i), cell)\n",
    "\n",
    "        for i in range(self.n_layers):\n",
    "            conv = nn.Sequential(nn.Conv2d(self.r_channels[i], self.a_channels[i], KERNEL_SIZE, padding=PADDING), nn.ReLU())\n",
    "            if i == 0:\n",
    "                conv.add_module('satlu', SatLU())\n",
    "            setattr(self, 'conv{}'.format(i), conv)\n",
    "\n",
    "        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "\n",
    "        for l in range(self.n_layers - 1):\n",
    "            update_A = nn.Sequential(nn.Conv2d(2* self.a_channels[l], self.a_channels[l+1], KERNEL_SIZE, padding=PADDING), self.maxpool)\n",
    "            setattr(self, 'update_A{}'.format(l), update_A)\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        for l in range(self.n_layers):\n",
    "            cell = getattr(self, 'cell{}'.format(l))\n",
    "            cell.reset_parameters()\n",
    "\n",
    "\n",
    "    def forward(self, input, time_steps=1, forecast_steps=0, mode='error'):\n",
    "        \n",
    "        R_seq = [None] * self.n_layers\n",
    "        H_seq = [None] * self.n_layers\n",
    "        E_seq = [None] * self.n_layers\n",
    "        \n",
    "        batch_size, input_steps, channels, width, height = input.size()\n",
    "        \n",
    "        w = width\n",
    "        h = height\n",
    "        for l in range(self.n_layers):\n",
    "            E_seq[l] = torch.randn(batch_size, 2*self.a_channels[l], w, h).cuda()\n",
    "            R_seq[l] = torch.randn(batch_size, self.r_channels[l], w, h).cuda()\n",
    "            w = w//2\n",
    "            h = h//2\n",
    "        \n",
    "        if mode=='error':\n",
    "            total_error = []\n",
    "        else:\n",
    "            output = []\n",
    "        \n",
    "        for t in range(time_steps + forecast_steps):\n",
    "            if t < input_steps:\n",
    "                frame_input = input[:,t]\n",
    "                frame_input = frame_input.type(torch.cuda.FloatTensor)\n",
    "            else:\n",
    "                frame_input = None\n",
    "            \n",
    "            for l in reversed(range(self.n_layers)):\n",
    "                cell = getattr(self, 'cell{}'.format(l))\n",
    "                \n",
    "                if t == 0:\n",
    "                    E = E_seq[l]\n",
    "                    R = R_seq[l]\n",
    "                    hx = (R, R)\n",
    "                else:\n",
    "                    E = E_seq[l]\n",
    "                    R = R_seq[l]\n",
    "                    hx = H_seq[l]\n",
    "                \n",
    "                if l == self.n_layers - 1:\n",
    "                    R, hx = cell(E, hx)\n",
    "                else:\n",
    "                    tmp = torch.cat((E, F.interpolate(R_seq[l+1], scale_factor=2)), 1)\n",
    "                    R, hx = cell(tmp, hx)\n",
    "                \n",
    "                R_seq[l] = R\n",
    "                H_seq[l] = hx\n",
    "                \n",
    "            for l in range(self.n_layers):\n",
    "                conv = getattr(self, 'conv{}'.format(l))\n",
    "                A_hat = conv(R_seq[l])\n",
    "                \n",
    "                if l == 0:\n",
    "                    frame_prediction = A_hat\n",
    "                    if t < time_steps:\n",
    "                        A = frame_input\n",
    "                    else:\n",
    "                        A = frame_prediction\n",
    "                \n",
    "                pos = F.relu(A_hat - A)\n",
    "                neg = F.relu(A - A_hat)\n",
    "                E = torch.cat([pos, neg],1)\n",
    "                E_seq[l] = E\n",
    "                \n",
    "                if l < self.n_layers - 1:\n",
    "                    update_A = getattr(self, 'update_A{}'.format(l))\n",
    "                    A = update_A(E)\n",
    "                \n",
    "            if mode == 'error':\n",
    "                error = torch.mean((frame_input - frame_prediction)**2)\n",
    "                total_error.append(error)\n",
    "            else:\n",
    "                output.append(frame_prediction)\n",
    "            \n",
    "        if mode == 'error':\n",
    "            return torch.stack(total_error, 0)\n",
    "        else:\n",
    "            return torch.stack(output, 1)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "iLCeS1oAYAEf",
    "outputId": "a7cc6d16-dfda-403e-ed9a-5b91cd1e0f95"
   },
   "outputs": [],
   "source": [
    "A_channels = (1, 16, 32, 64)\n",
    "#A_channels = (1, 32, 64, 128)\n",
    "R_channels = (1, 16, 32, 64)\n",
    "#R_channels = (1, 32, 64, 128)\n",
    "\n",
    "model = PredNet(R_channels, A_channels)\n",
    "if torch.cuda.is_available():\n",
    "    print('Using GPU.')\n",
    "    model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "jp4_2I8gYAEi",
    "outputId": "a642c0e0-ff96-4644-f5dd-823d37cd4a32"
   },
   "outputs": [],
   "source": [
    "checkpoint = torch.load('/content/gdrive/My Drive/Nowcasting/KyotoRadar/prednet-v4.pth')\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "#optimizer.load_state_dict(checkpoint['optimizer_state_dict'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 765
    },
    "colab_type": "code",
    "id": "sKBwutQOYAEl",
    "outputId": "5eb5d36b-7c67-40bc-f688-5acb7e763cdc"
   },
   "outputs": [],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "LC0iZ8GHYAEo",
    "outputId": "3d2b1fd4-2f83-4d1f-d29e-31f57ae4096d"
   },
   "outputs": [],
   "source": [
    "pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print(pytorch_total_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "RKU_IVifYAEr",
    "outputId": "8f75d9e2-7e33-483b-9cf8-e801b1a121b2"
   },
   "outputs": [],
   "source": [
    "e = model(images, time_steps=10, mode='error')\n",
    "print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "z76CpzN9YAEw",
    "outputId": "69f56f09-4f3d-4968-e0c2-bfc2f8228f13"
   },
   "outputs": [],
   "source": [
    "o = model(images, time_steps=10, forecast_steps=5, mode='forecast')\n",
    "print(o.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tV3V1hBfYAE0"
   },
   "outputs": [],
   "source": [
    "train_losses = []\n",
    "train_counter = []\n",
    "valid_losses = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jqd27GmJYAE1"
   },
   "outputs": [],
   "source": [
    "def train(epoch, log_interval):\n",
    "    \"\"\"Training\"\"\"\n",
    "    model.train()\n",
    "    loss_log = 0\n",
    "    for batch_idx, frames in enumerate(train_loader):\n",
    "        errors = model(frames, time_steps=12, forecast_steps=12, mode='error')\n",
    "        loss = torch.mean(errors)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        loss_log += loss.item()\n",
    "\n",
    "        if ((batch_idx + 1) % log_interval == 0):\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch, (batch_idx +1) * len(frames), len(train_loader)*BATCH_SIZE,\n",
    "                100. * (batch_idx +1) / len(train_loader), loss_log / log_interval))\n",
    "            loss_log = 0\n",
    "            train_losses.append(loss.item())\n",
    "            train_counter.append((batch_idx*BATCH_SIZE) + ((epoch-1)*len(train_loader)*BATCH_SIZE))\n",
    "            #print(' > Epoch {:2d} loss: {:.3f}'.format((epoch + 1), loss.data))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "c-V3onYRYAE5"
   },
   "outputs": [],
   "source": [
    "def valid():\n",
    "    \"\"\"Validation\"\"\"\n",
    "    model.eval()\n",
    "    loss_log = 0\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, frames in enumerate(valid_loader):\n",
    "            errors = model(frames, time_steps=12, forecast_steps=12, mode='error')\n",
    "            loss = torch.mean(errors)\n",
    "\n",
    "            loss_log += loss.item()\n",
    "        \n",
    "        loss_log /= len(valid_loader)\n",
    "        valid_losses.append(loss_log)\n",
    "        print('\\nValidation loss: {:.6f}\\n'.format(loss_log))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "oUus_dWhYAE6",
    "outputId": "2ef58c8f-bd34-4c6a-9611-e6460d800293"
   },
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()\n",
    "print(torch.cuda.memory_allocated()/(1024**2), torch.cuda.max_memory_allocated()/(1024**2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qaSjojyxYAE8"
   },
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "H0pbyNKBTnVK"
   },
   "outputs": [],
   "source": [
    "#checkpoint = torch.load('/content/gdrive/My Drive/Nowcasting/KyotoRadar/prednet-out.pth')\n",
    "#model.load_state_dict(checkpoint['model_state_dict'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 531
    },
    "colab_type": "code",
    "id": "qVpCbmbzYAFA",
    "outputId": "2b12bcf3-e0a5-4e71-dd4c-09257fd2733b"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "log_interval = 10\n",
    "n_epochs = 4\n",
    "\n",
    "for epoch in range(1, n_epochs + 1):\n",
    "    train(epoch, log_interval)\n",
    "    valid()\n",
    "    \n",
    "    torch.save({\n",
    "        'epoch': epoch,\n",
    "        'model_state_dict': model.state_dict(),\n",
    "        'optimizer_state_dict': optimizer.state_dict()\n",
    "        }, '/content/gdrive/My Drive/Nowcasting/KyotoRadar/prednet-out.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YDwDzWm0Ups6"
   },
   "outputs": [],
   "source": [
    "torch.save({\n",
    "    'epoch': epoch,\n",
    "    'model_state_dict': model.state_dict(),\n",
    "    'optimizer_state_dict': optimizer.state_dict()\n",
    "    }, '/content/gdrive/My Drive/Nowcasting/KyotoRadar/prednet-out.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "WRVMziVKYAFC",
    "outputId": "9145c56d-367f-44bc-f387-2eacb80d9be8"
   },
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()\n",
    "print(torch.cuda.memory_allocated()/(1024**2), torch.cuda.max_memory_allocated()/(1024**2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aIJP02L26ZeM"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "WYJPf5ksYAFE",
    "outputId": "02abab08-3663-4d2b-c085-7b7532895671"
   },
   "outputs": [],
   "source": [
    "checkpoint = torch.load('/content/gdrive/My Drive/Nowcasting/KyotoRadar/prednet-All.pth')\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "#optimizer.load_state_dict(checkpoint['optimizer_state_dict'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "L1grIYngYAFI",
    "outputId": "8c108053-e095-44ff-9e9c-43e894262d36"
   },
   "outputs": [],
   "source": [
    "images.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 297
    },
    "colab_type": "code",
    "id": "8rI0PiHpYAFO",
    "outputId": "bc7d90d9-3b85-4f52-ead9-353f55c5a5b5"
   },
   "outputs": [],
   "source": [
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    idx = 26 #894 \n",
    "    print(valid_times[idx])\n",
    "    input_sequence = torch.from_numpy(radardata_valid[idx]).unsqueeze(0)\n",
    "    #input_sequence = images.cuda()\n",
    "    \n",
    "    plt.figure(figsize=(18,4))\n",
    "    b = 0\n",
    "    for i in range(12):\n",
    "        plt.subplot(4,12,i+1)\n",
    "        plt.tight_layout()\n",
    "        plt.imshow(input_sequence[0,i,0].cpu().detach().numpy(), vmin=0, vmax=5/15)\n",
    "        #plt.imshow(input_sequence[b,i,0].cpu().detach().numpy())\n",
    "        plt.axis('off')\n",
    "        \n",
    "        plt.subplot(4,12,i+25)\n",
    "        plt.tight_layout()\n",
    "        plt.imshow(input_sequence[0,i+12,0].cpu().detach().numpy(), vmin=0, vmax=5/15)\n",
    "        #plt.imshow(input_sequence[b,i+12,0].cpu().detach().numpy())\n",
    "        plt.axis('off')\n",
    "\n",
    "    output_sequence = model(input_sequence[:,:12,...], time_steps=12, forecast_steps=12, mode='forecast')\n",
    "\n",
    "    for i in range(12):\n",
    "        plt.subplot(4,12,i+13)\n",
    "        plt.tight_layout()\n",
    "        plt.imshow(output_sequence[b,i,0].cpu().detach().numpy(), vmin=0, vmax=5/15)\n",
    "        plt.axis('off')\n",
    "\n",
    "        plt.subplot(4,12,i+37)\n",
    "        plt.tight_layout()\n",
    "        plt.imshow(output_sequence[b,i+12,0].cpu().detach().numpy(), vmin=0, vmax=5/15)\n",
    "        plt.axis('off')\n",
    "    \n",
    "    rmse = np.sqrt(((input_sequence[0,12:,0]*15 - output_sequence[0,12:,0].cpu().detach()*15)**2).mean())\n",
    "    print(rmse.item())\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HVCoJxprgvD5"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "blr6WtkNpE9n"
   },
   "outputs": [],
   "source": [
    "def class2rainrate(rr):\n",
    "    amount = [0,1,2,4,8,12,16,24,32,40,48,56,64,80,100]\n",
    "    rt=rr.copy().astype(float)\n",
    "    for i,r in enumerate(amount):\n",
    "        rt[rr==i]=amount[i]\n",
    "    return rt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iF59nKAM_XM_"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "RainPredNet.ipynb",
   "provenance": []
  },
  "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}