ARCNet-Small-PAMAP.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZfYUd0VTCb-p"
   },
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iPR8tPoKRcKl"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scipy.io as io\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from dataclasses import dataclass, field, asdict\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import pytorch_lightning as pl\n",
    "from torch.nn import init\n",
    "from torch.optim import lr_scheduler\n",
    "from pytorch_lightning import Trainer\n",
    "from torch.autograd import Variable, Function\n",
    "from pytorch_lightning.callbacks import EarlyStopping\n",
    "from torch.nn.init import kaiming_normal_, orthogonal_\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from sklearn import metrics\n",
    "import json\n",
    "import torch\n",
    "import math\n",
    "import time\n",
    "import warnings\n",
    "import os\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RQYykKS6R-I8"
   },
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class paramclass:\n",
    "    lr = 1e-03\n",
    "    l = 0.1\n",
    "    batch_size: int = 64\n",
    "    routing_iterations: int = 5\n",
    "    n_classes: int = 12\n",
    "    max_epochs = 300\n",
    "    device = torch.device('cuda:0')\n",
    "    num_workers: int = 0\n",
    "    first_training = True\n",
    "    resume: str = ''\n",
    "\n",
    "params = paramclass()\n",
    "params.data_path = './data/pamap_3x6.mat'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ji4m-UCfRf0G"
   },
   "outputs": [],
   "source": [
    "# Preprocessing from https://arxiv.org/abs/1811.00170\n",
    "class dataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, file_paths, params, type='train', device=torch.device('cuda:0')):\n",
    "        self.params = params\n",
    "        self.device = device\n",
    "\n",
    "        pamap = io.loadmat(self.params.data_path)\n",
    "        X_train = np.asarray(pamap['train_X'][:,:,:]).reshape([pamap['train_X'].shape[0],18,250,1]) \n",
    "        y_train = np.asarray(pamap['train_y'].T, dtype=int) - 1\n",
    "        X_test = np.asarray(pamap['test_X'][:,:,:]).reshape([pamap['test_X'].shape[0],18,250,1])\n",
    "        y_test = np.asarray(pamap['test_y'].T, dtype=int) - 1\n",
    "        X_val = np.asarray(pamap['val_X'][:,:,:]).reshape([pamap['val_X'].shape[0],18,250,1])\n",
    "        y_val = np.asarray(pamap['val_y'].T, dtype=int) - 1\n",
    "            \n",
    "        for i in range(18):\n",
    "            train_mean = np.mean(X_train[:,i,:,:])\n",
    "            train_std = np.std(X_train[:,i,:,:])\n",
    "            X_train[:,i,:,:] = (X_train[:,i,:,:] - train_mean)/train_std\n",
    "\n",
    "            test_mean = np.mean(X_test[:,i,:,:])\n",
    "            test_std = np.std(X_test[:,i,:,:])\n",
    "            X_test[:,i,:,:] = (X_test[:,i,:,:] - test_mean)/test_std\n",
    "                \n",
    "            val_mean = np.mean(X_val[:,i,:,:])\n",
    "            val_std = np.std(X_val[:,i,:,:])\n",
    "            X_val[:,i,:,:] = (X_val[:,i,:,:] - val_mean)/val_std\n",
    "        \n",
    "        X_train = X_train.reshape([13720,18,250,1])[:,:,:128,:]\n",
    "        X_test = X_test.reshape([2405,18,250,1])[:,:,:128,:]\n",
    "        X_val = X_val.reshape([2640,18,250,1])[:,:,:128,:]\n",
    "        X_train = X_train[:, :, :, 0]\n",
    "        X_test = X_test[:, :, :, 0]\n",
    "        X_val = X_val[:, :, :, 0]\n",
    "\n",
    "        y_train_onehot = np.zeros((y_train.shape[0], 12), np.uint8)\n",
    "        y_val_onehot = np.zeros((y_val.shape[0], 12), np.uint8)\n",
    "        y_test_onehot = np.zeros((y_test.shape[0], 12), np.uint8)\n",
    "        y_train_onehot[np.arange(y_train.shape[0]).squeeze(), y_train.squeeze()] = 1\n",
    "        y_val_onehot[np.arange(y_val.shape[0]).squeeze(), y_val.squeeze()] = 1\n",
    "        y_test_onehot[np.arange(y_test.shape[0]).squeeze(), y_test.squeeze()] = 1\n",
    "\n",
    "        if type=='train':\n",
    "            self.X = X_train\n",
    "            self.y = y_train_onehot\n",
    "        elif type=='test':\n",
    "            self.X = X_test\n",
    "            self.y = y_test_onehot\n",
    "        else:\n",
    "            self.X = X_val\n",
    "            self.y = y_val_onehot\n",
    "        self.datalen = self.X.shape[0]\n",
    "\n",
    "    def preprocess(self, data):\n",
    "        B, H, W = data.shape\n",
    "        data = data.view(B, 1, H, W)\n",
    "        imu1 = data[:, :,   :6,   :]\n",
    "        imu2 = data[:, :,  6:12,  :]\n",
    "        imu3 = data[:, :, 12:,    :]\n",
    "        return imu1.float(), imu2.float(), imu3.float()\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.datalen\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "        data = self.X[idx]\n",
    "        target = self.y[idx]\n",
    "        return data, target\n",
    "\n",
    "def conv2d(batchNorm, in_channels, out_channels, kernel_size, stride=1):\n",
    "    if batchNorm:\n",
    "        return nn.Sequential(\n",
    "            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=False),\n",
    "            nn.BatchNorm2d(out_channels),\n",
    "            nn.LeakyReLU(0.2,inplace=True),\n",
    "            nn.Dropout2d(p=0.35)\n",
    "        )\n",
    "    else:\n",
    "        return nn.Sequential(\n",
    "            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=True),\n",
    "            nn.LeakyReLU(0.2,inplace=True),\n",
    "            nn.Dropout2d(p=0.35)\n",
    "        )\n",
    "\n",
    "def squash(x):\n",
    "    lengths2 = x.pow(2).sum(dim=2)\n",
    "    lengths = lengths2.sqrt()\n",
    "    x = x * (lengths2 / (1 + lengths2) / lengths).view(x.size(0), x.size(1), 1)\n",
    "    return x\n",
    "\n",
    "class MarginLoss(pl.LightningModule):\n",
    "    def __init__(self, m_pos, m_neg, lambda_):\n",
    "        super(MarginLoss, self).__init__()\n",
    "        self.m_pos = m_pos\n",
    "        self.m_neg = m_neg\n",
    "        self.lambda_ = lambda_\n",
    "\n",
    "    def forward(self, lengths, targets, size_average=True):\n",
    "        losses = targets.float() * F.relu(self.m_pos - lengths).pow(2) + \\\n",
    "                 self.lambda_ * (1. - targets.float()) * F.relu(lengths - self.m_neg).pow(2)\n",
    "        return losses.mean() if size_average else losses.sum()\n",
    "\n",
    "class Small_Encoder(pl.LightningModule):\n",
    "    def __init__(self, batchNorm):\n",
    "        super(Small_Encoder,self).__init__()\n",
    "\n",
    "        self.batchNorm = batchNorm\n",
    "\n",
    "        self.conv1   = conv2d(self.batchNorm,   1,   16, kernel_size=(1, 10), stride=(1, 2))\n",
    "        self.conv2   = conv2d(self.batchNorm,  16,  32, kernel_size=(3, 10), stride=(3, 2))\n",
    "        self.conv3   = conv2d(self.batchNorm, 32,  48, kernel_size=(2, 15), stride=1)\n",
    "        \n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):\n",
    "                kaiming_normal_(m.weight.data)\n",
    "                if m.bias is not None:\n",
    "                    m.bias.data.zero_()\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                m.weight.data.fill_(1)\n",
    "                m.bias.data.zero_()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.conv3(x)\n",
    "        return x\n",
    "    \n",
    "class AgreementRouting(pl.LightningModule):\n",
    "    def __init__(self, input_caps, output_caps, n_iterations, l):\n",
    "        super(AgreementRouting, self).__init__()\n",
    "        self.n_iterations = n_iterations\n",
    "        self.b = nn.Parameter(torch.zeros((input_caps, output_caps)))\n",
    "        self.l = l\n",
    "\n",
    "    def forward(self, u_predict):\n",
    "        batch_size, input_caps, output_caps, output_dim = u_predict.size()\n",
    "\n",
    "        c = F.softmax(self.b)\n",
    "        s = (c.unsqueeze(2) * u_predict).sum(dim=1)\n",
    "        v = squash(s)\n",
    "\n",
    "        if self.n_iterations > 0:\n",
    "            b_batch = self.b.expand((batch_size, input_caps, output_caps))\n",
    "            for r in range(self.n_iterations):\n",
    "                v = v.unsqueeze(1)\n",
    "                b_batch = (1-self.l)*b_batch + self.l*(u_predict * v).sum(-1)\n",
    "\n",
    "                c = F.softmax(b_batch.view(-1, output_caps)).view(-1, input_caps, output_caps, 1)\n",
    "                s = (c * u_predict).sum(dim=1)\n",
    "                v = squash(s)\n",
    "        return v\n",
    "\n",
    "    \n",
    "class CapsLayer(pl.LightningModule):\n",
    "    def __init__(self, input_caps, input_dim, output_caps, output_dim, routing_module):\n",
    "        super(CapsLayer, self).__init__()\n",
    "        self.input_dim = input_dim\n",
    "        self.input_caps = input_caps\n",
    "        self.output_dim = output_dim\n",
    "        self.output_caps = output_caps\n",
    "        self.weights = nn.Parameter(torch.Tensor(input_caps, input_dim, output_caps * output_dim))\n",
    "        self.routing_module = routing_module\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        stdv = 1. / math.sqrt(self.input_caps)\n",
    "        self.weights.data.uniform_(-stdv, stdv)\n",
    "\n",
    "    def forward(self, caps_output):\n",
    "        caps_output = caps_output.unsqueeze(2)\n",
    "        u_predict = caps_output.matmul(self.weights)\n",
    "        u_predict = u_predict.view(u_predict.size(0), self.input_caps, self.output_caps, self.output_dim)\n",
    "        v = self.routing_module(u_predict)\n",
    "        return v\n",
    "\n",
    "    \n",
    "class PrimaryCapsLayer(pl.LightningModule):\n",
    "    def __init__(self, split=False, output_caps=None, output_dim=None):\n",
    "        super(PrimaryCapsLayer, self).__init__()\n",
    "        self.split = split\n",
    "        self.output_caps = output_caps\n",
    "        self.output_dim = output_dim\n",
    "    def forward(self, input):\n",
    "        [i1, i2, i3] = input\n",
    "        if self.split:\n",
    "            B, C, H, W = i1.shape\n",
    "            i1 = i1.view(B, self.output_caps, self.output_dim, H, W)\n",
    "            i2 = i2.view(B, self.output_caps, self.output_dim, H, W)\n",
    "            i3 = i3.view(B, self.output_caps, self.output_dim, H, W)\n",
    "            i1 = i1.permute(0, 1, 3, 4, 2).contiguous()\n",
    "            i2 = i2.permute(0, 1, 3, 4, 2).contiguous()\n",
    "            i3 = i3.permute(0, 1, 3, 4, 2).contiguous()\n",
    "            i1 = i1.view(i1.size(0), -1, i1.size(4))\n",
    "            i2 = i2.view(i2.size(0), -1, i2.size(4))\n",
    "            i3 = i3.view(i3.size(0), -1, i3.size(4))\n",
    "            out = torch.cat([i1, i2, i3], dim=1)\n",
    "        else:\n",
    "            B, C, H, W = i1.shape\n",
    "            i1 = i1.permute(0, 2, 3, 1).contiguous()\n",
    "            i2 = i2.permute(0, 2, 3, 1).contiguous()\n",
    "            i3 = i3.permute(0, 2, 3, 1).contiguous()\n",
    "            i1 = i1.view(B, H*W, C)\n",
    "            i2 = i2.view(B, H*W, C)\n",
    "            i3 = i3.view(B, H*W, C)\n",
    "            out = torch.cat([i1, i2, i3], dim=1)\n",
    "        out = squash(out)\n",
    "        return out\n",
    "\n",
    "class CapsNet(pl.LightningModule):\n",
    "    def __init__(self, hparams):\n",
    "        super(CapsNet, self).__init__()\n",
    "        if type(hparams)==dict:\n",
    "            self.hparams = hparams\n",
    "            self.params = paramclass(**hparams)\n",
    "        else:\n",
    "            self.hparams = asdict(hparams)\n",
    "            self.params = hparams\n",
    "\n",
    "        self.pp = Small_Encoder(False)\n",
    "        self.primaryCaps = PrimaryCapsLayer()\n",
    "        self.num_primaryCaps = 1*12*3\n",
    "        routing_module = AgreementRouting(self.num_primaryCaps, self.params.n_classes, self.params.routing_iterations, self.params.l)\n",
    "        self.activityCaps = CapsLayer(self.num_primaryCaps, 48, self.params.n_classes, 16, routing_module)\n",
    "            \n",
    "        if not self.params.first_training:\n",
    "            self.load_from_checkpoint(self.params.resume)\n",
    "\n",
    "        self.loss_fn = MarginLoss(0.95, 0.05, 0.5)\n",
    "        \n",
    "    def prepare_data(self):\n",
    "        self.traindataArr = dataset(self.params.data_path, self.params, type='train')\n",
    "        self.valdataArr = dataset(self.params.data_path, self.params, type='val')\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(self.traindataArr, batch_size=self.params.batch_size, shuffle=True, num_workers=self.params.num_workers, drop_last=True)\n",
    "    \n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(self.valdataArr, batch_size=1, shuffle=False, num_workers=0)\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        optimizer = optim.Adam(self.parameters(), lr=self.params.lr)\n",
    "        scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.986)\n",
    "        return [optimizer], [scheduler]\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        data, target = batch\n",
    "        imu1, imu2, imu3 = self.traindataArr.preprocess(data)\n",
    "        output, probs = self([imu1, imu2, imu3])\n",
    "        loss = self.loss_fn(probs, target)\n",
    "        tensorboard_logs = {'train_loss': loss}\n",
    "        return {'loss': loss, 'log': tensorboard_logs}\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        data, target = batch\n",
    "        imu1, imu2, imu3 = self.traindataArr.preprocess(data)\n",
    "        output, probs = self([imu1, imu2, imu3])\n",
    "        loss = self.loss_fn(probs, target)\n",
    "        return {'val_loss': loss.data, 'Target': target, 'Predictions': probs}\n",
    "    \n",
    "    def validation_epoch_end(self, outputs):\n",
    "        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()\n",
    "        targets = torch.stack([x['Target'] for x in outputs])[:, 0, :].cpu().numpy()\n",
    "        predictions = torch.stack([x['Predictions'] for x in outputs])[:, 0, :].cpu().numpy() \n",
    "        temp = np.where(predictions==predictions.max(axis=1, keepdims=True))[1]\n",
    "        hotpredictions = np.zeros((temp.size, 12))\n",
    "        hotpredictions[np.arange(temp.size), temp] = 1\n",
    "        f1_mac = metrics.f1_score(targets, hotpredictions, average='macro')\n",
    "        f1_mic = metrics.f1_score(targets, hotpredictions, average='micro')\n",
    "        acc = metrics.accuracy_score(targets, hotpredictions)\n",
    "        precision_avg, recall_avg, f_score_avg,_avg = metrics.precision_recall_fscore_support(targets, hotpredictions, average='weighted')\n",
    "        tensorboard_logs = {'avg_val_loss': avg_loss,\n",
    "                            'F1-Micro': f1_mic,\n",
    "                            'F1-Macro': f1_mac,\n",
    "                            'Accuracy': acc,\n",
    "                            'Precision_Avg': precision_avg,\n",
    "                            'Recall_Avg': recall_avg,\n",
    "                            'F_score_Avg': f_score_avg\n",
    "                            }\n",
    "        return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}\n",
    "                \n",
    "    def forward(self, input):\n",
    "        capsules = []\n",
    "        for imu in input:\n",
    "            capsules.append(self.pp(imu))\n",
    "        x = self.primaryCaps(capsules)\n",
    "        x = self.activityCaps(x)\n",
    "        probs = x.pow(2).sum(dim=2).sqrt()\n",
    "        return x, probs\n",
    "    \n",
    "class Non_val_epoch_saves(pl.Callback):\n",
    "    def __init__(self, iteration, filepath):\n",
    "        self.iteration = iteration\n",
    "        self.filepath = filepath    \n",
    "    def on_epoch_end(self, trainer, pl_module):\n",
    "        self.name = self.iteration + '_Epoch=' + str(trainer.current_epoch) + '.ckpt'\n",
    "        trainer.checkpoint_callback._save_model(filepath=os.path.join(self.filepath, self.name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 442
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 1745423,
     "status": "ok",
     "timestamp": 1591449354927,
     "user": {
      "displayName": "Hamed Damirchi",
      "photoUrl": "",
      "userId": "01846405488099323235"
     },
     "user_tz": -270
    },
    "id": "blPn7fh8yrBa",
    "outputId": "e2088890-5035-42d1-f4d4-4e45e6d3ce84"
   },
   "outputs": [],
   "source": [
    "model = CapsNet(params)\n",
    "iteration = 'ARCNet-Small-PAMAP'\n",
    "callback_dir = './checkpoints/'\n",
    "\n",
    "trainer = Trainer(gpus=1,\n",
    "                  default_save_path='./checkpoints/',\n",
    "                  track_grad_norm=2,\n",
    "                  max_epochs=params.max_epochs,\n",
    "                  progress_bar_refresh_rate=50,\n",
    "                  weights_summary='top',\n",
    "                  callbacks=[Non_val_epoch_saves(iteration=iteration, filepath=callback_dir)])\n",
    "trainer.fit(model)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "M10-3.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python (torch14)",
   "language": "python",
   "name": "torch14"
  },
  "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.7.3"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "384px"
   },
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}