4_LSTM.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Long short-term memory\n",
    "\n",
    "We saw that accounting for the temporal dependency is critical to improve\n",
    "performance.  We found improvements with a hidden Markov model, but also using a\n",
    "simple mode smoothing.\n",
    "Here we look at using a more flexible model &mdash; the Long short-term memory\n",
    "(LSTM) &mdash; to model the temporal dependency and smooth the predictions of a\n",
    "random forest.\n",
    "\n",
    "## AWS info\n",
    "Run this notebook with the **PyTorch 1.12 Python 3.8 GPU Optimized** image and **ml.g4dn.2xlarge** instance type.\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from imblearn.ensemble import BalancedRandomForestClassifier\n",
    "from sklearn import metrics\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.backends.cudnn as cudnn\n",
    "\n",
    "import utils\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "%matplotlib inline\n",
    "\n",
    "# For reproducibility\n",
    "np.random.seed(42)\n",
    "torch.manual_seed(42)\n",
    "cudnn.benchmark = True\n",
    "\n",
    "# Grab a GPU if there is one\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    print(\"Using {} device: {}\".format(device, torch.cuda.current_device()))\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    print(\"Using {}\".format(device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path to your extracted windows\n",
    "DATASET_PATH = 'processed_data/'\n",
    "X_FEATS_PATH = 'X_feats.pkl'  # path to your extracted features, if have one\n",
    "print(f'Content of {DATASET_PATH}')\n",
    "print(os.listdir(DATASET_PATH))\n",
    "\n",
    "X = np.load(DATASET_PATH+'X.npy')\n",
    "Y = np.load(DATASET_PATH+'Y.npy')\n",
    "T = np.load(DATASET_PATH+'T.npy')\n",
    "pid = np.load(DATASET_PATH+'pid.npy')\n",
    "X_feats = pd.read_pickle(DATASET_PATH+'X_feats.pkl')\n",
    "\n",
    "# As before, let's map the text annotations to simplified labels\n",
    "ANNO_LABEL_DICT_PATH = 'capture24/annotation-label-dictionary.csv'\n",
    "anno_label_dict = pd.read_csv(ANNO_LABEL_DICT_PATH, index_col='annotation', dtype='string')\n",
    "Y = anno_label_dict.loc[Y, 'label:Willetts2018'].to_numpy()\n",
    "\n",
    "# Transform to numeric\n",
    "le = LabelEncoder().fit(Y)\n",
    "Y = le.transform(Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train/test split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hold out participants P101-P151 for testing (51 participants)\n",
    "test_ids = [f'P{i}' for i in range(101,152)]\n",
    "mask_test = np.isin(pid, test_ids)\n",
    "mask_train = ~mask_test\n",
    "X_train, Y_train, T_train, pid_train = \\\n",
    "    X_feats[mask_train], Y[mask_train], T[mask_train], pid[mask_train]\n",
    "X_test, Y_test, T_test, pid_test = \\\n",
    "    X_feats[mask_test], Y[mask_test], T[mask_test], pid[mask_test]\n",
    "print(\"Shape of X_train:\", X_train.shape)\n",
    "print(\"Shape of X_test:\", X_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train a random forest classifier\n",
    "\n",
    "*Note: this may take a while*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clf = BalancedRandomForestClassifier(\n",
    "    n_estimators=1000,\n",
    "    replacement=True,\n",
    "    sampling_strategy='not minority',\n",
    "    oob_score=True,\n",
    "    n_jobs=-1,\n",
    "    random_state=42,\n",
    "    verbose=1\n",
    ")\n",
    "clf.fit(X_train, Y_train)\n",
    "\n",
    "Y_test_pred = clf.predict(X_test)\n",
    "print('\\nClassifier performance')\n",
    "print('Out of sample:\\n', metrics.classification_report(Y_test, Y_test_pred, zero_division=0))\n",
    "\n",
    "# This will be the training set\n",
    "Y_in_train = clf.oob_decision_function_.astype('float32')\n",
    "# This will be the test set\n",
    "Y_in_test = clf.predict_proba(X_test).astype('float32')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Architecture design\n",
    "As a baseline, let's use a single-layer bidirectional LSTM.\n",
    "PyTorch uses a sligtly unintuitive array format for the input and output of\n",
    "its LSTM module.\n",
    "The array shape for both input and output is `(seq_length,N,num_labels)`, corresponding to\n",
    "`N` sequences of `seq_length` elements of size `num_labels`.\n",
    "Here, each element is a vector of label probabilities/logits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTM(nn.Module):\n",
    "    ''' Single-layer bidirectional LSTM '''\n",
    "    def __init__(self, input_size=5, output_size=5, hidden_size=1024, num_layers=2, bias=True, dropout=.5):\n",
    "        super(LSTM, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.output_size = output_size\n",
    "        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, bias, dropout=dropout, bidirectional=True)\n",
    "        self.hidden2output = nn.Linear(2*hidden_size, output_size)\n",
    "\n",
    "    def forward(self, sequence):\n",
    "        hiddens, (hidden_last, cell_last) = self.lstm(\n",
    "            sequence.view(len(sequence), -1, self.input_size))\n",
    "        output = self.hidden2output(\n",
    "            hiddens.view(-1, hiddens.shape[-1])).view(\n",
    "                hiddens.shape[0], hiddens.shape[1], self.output_size\n",
    "        )\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataloader(Y, y=None, seq_length=5, batch_size=1, shuffle=False, eval_mode=False):\n",
    "    ''' Create a (batch) iterator over the dataset. It yields (batches of)\n",
    "    sequences of consecutive rows of `Y` and `y` of length `seq_length` (can\n",
    "    be less than `seq_length` in `eval_mode=True`).\n",
    "\n",
    "    The below code looks complicated but all it's trying to do is to pack\n",
    "    sequences of equal length where applicable, else provide the sequences\n",
    "    one by one.\n",
    "    '''\n",
    "    if eval_mode:\n",
    "        # In order to reuse this loader in evaluation/prediction mode, we\n",
    "        # provide non-overlapping segments, as well as the trailing segments\n",
    "        # that can be shorter than seq_length.\n",
    "        n = len(Y)\n",
    "        idxs = np.arange(0, n, seq_length)\n",
    "    else:\n",
    "        n = len(Y) - seq_length + 1\n",
    "        idxs = np.arange(n)\n",
    "    if shuffle:\n",
    "        idxs = np.random.permutation(idxs)\n",
    "    for i in range(0, len(idxs), batch_size):\n",
    "        idxs_batch = idxs[i:i+batch_size]\n",
    "        # Separate those with irregular length -- these will be yielded one by one\n",
    "        idxs_batch_regular = np.asarray(\n",
    "            [j for j in idxs_batch if len(Y[j:j+seq_length]) == seq_length]\n",
    "        )\n",
    "        idxs_batch_irregular = np.asarray(\n",
    "            [j for j in idxs_batch if j not in idxs_batch_regular]\n",
    "        )\n",
    "        # Yield batch of sequences of regular length (=seq_length)\n",
    "        sequence_batch = np.stack([Y[j:j+seq_length] for j in idxs_batch_regular], axis=1)\n",
    "        sequence_batch = torch.from_numpy(sequence_batch)\n",
    "        if y is None:\n",
    "            yield sequence_batch\n",
    "        else:\n",
    "            y_batch = np.stack([y[j:j+seq_length] for j in idxs_batch_regular], axis=1)\n",
    "            y_batch = torch.from_numpy(y_batch)\n",
    "            yield sequence_batch, y_batch\n",
    "        # Yield sequences of irregular length uno por uno\n",
    "        for j in idxs_batch_irregular:\n",
    "            sequence_batch = torch.from_numpy(Y[j:j+seq_length]).unsqueeze(1)\n",
    "            if y is None:\n",
    "                yield sequence_batch\n",
    "            else:\n",
    "                y_batch = torch.from_numpy(y[j:j+seq_length]).unsqueeze(1)\n",
    "                yield sequence_batch, y_batch\n",
    "\n",
    "\n",
    "def forward_by_batches(lstm, Y_in, seq_length):\n",
    "    ''' Forward pass model on a dataset.\n",
    "    Do this by batches so that we don't blow up the memory. '''\n",
    "    Y_out = []\n",
    "    lstm.eval()\n",
    "    with torch.no_grad():\n",
    "        for sequence in create_dataloader(\n",
    "            Y_in, seq_length=seq_length, batch_size=1024, shuffle=False, eval_mode=True\n",
    "        ):  # do not shuffle here!\n",
    "            sequence = sequence.to(device)\n",
    "            output = lstm(sequence)\n",
    "            Y_out.append(output)\n",
    "    lstm.train()\n",
    "    # Concatenate sequences in order -- need to transpose to get batch-first format\n",
    "    Y_out = torch.cat(\n",
    "        [output.transpose(1,0).reshape(-1, output.shape[-1]) for output in Y_out]\n",
    "    )\n",
    "    return Y_out\n",
    "\n",
    "\n",
    "def evaluate_model(lstm, Y_in, Y, seq_length):\n",
    "    Y_pred_prob = forward_by_batches(lstm, Y_in, seq_length)  # lstm smoothing (scores)\n",
    "    loss = F.cross_entropy(Y_pred_prob, torch.from_numpy(Y).type(torch.int64).to(device)).item()\n",
    "\n",
    "    Y_pred_prob = F.softmax(Y_pred_prob, dim=-1)  # convert to probabilities\n",
    "    Y_pred = torch.argmax(Y_pred_prob, dim=-1)  # convert to classes\n",
    "    Y_pred = Y_pred.cpu().numpy()  # cast to numpy array\n",
    "    kappa = metrics.cohen_kappa_score(Y, Y_pred)\n",
    "\n",
    "    return {'loss':loss, 'kappa':kappa, 'Y_pred':Y_pred}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## Hyperparameters, model instantiation, loss function and optimizer "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_size = 1024  # size of LSTM's hidden state\n",
    "num_layers = 2  # num layers in the LSTM module\n",
    "dropout = .5  # dropout rate in LSTM module\n",
    "input_size = output_size = len(np.unique(Y))\n",
    "seq_length = 5  # max num of elems to consider for smoothing (the time horizon)\n",
    "num_epoch = 5  # num of epochs (full loops though the training set)\n",
    "lr = 3e-4  # learning rate\n",
    "batch_size = 32  # size of the mini-batch\n",
    "\n",
    "lstm = LSTM(\n",
    "    input_size=input_size,\n",
    "    output_size=output_size,\n",
    "    hidden_size=hidden_size,\n",
    "    num_layers=num_layers,\n",
    "    dropout=dropout\n",
    ").to(device)\n",
    "print(lstm)\n",
    "\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(lstm.parameters(), lr=lr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## Training "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kappa_history_test = []\n",
    "loss_history_test = []\n",
    "loss_history_train = []\n",
    "losses = []\n",
    "for i in tqdm(range(num_epoch)):\n",
    "    dataloader = create_dataloader(Y_in_train, Y_train, seq_length, batch_size, shuffle=True)\n",
    "    for sequence, target in dataloader:\n",
    "        sequence, target = sequence.to(device), target.type(torch.int64).to(device)\n",
    "        lstm.zero_grad()\n",
    "        output = lstm(sequence)\n",
    "        loss = loss_fn(output.view(-1,output.shape[-1]), target.view(-1))\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # Logging -- track train loss\n",
    "        losses.append(loss.item())\n",
    "\n",
    "    # --------------------------------------------------------\n",
    "    #       Evaluate performance at the end of each epoch\n",
    "    # --------------------------------------------------------\n",
    "\n",
    "    # Logging -- average train loss in this epoch\n",
    "    loss_history_train.append(utils.ewm(losses))\n",
    "\n",
    "    # Logging -- evaluate performance on test set\n",
    "    results = evaluate_model(lstm, Y_in_test, Y_test, seq_length)\n",
    "    loss_history_test.append(results['loss'])\n",
    "    kappa_history_test.append(results['kappa'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## Model performane "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loss history\n",
    "plt.close('all')\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(loss_history_train, color='C0', label='train loss')\n",
    "ax.plot(loss_history_test, color='C1', label='test loss')\n",
    "ax.set_ylabel('loss (CE)')\n",
    "ax.set_xlabel('epoch')\n",
    "ax = ax.twinx()\n",
    "ax.plot(kappa_history_test, color='C2', label='kappa')\n",
    "ax.set_ylabel('kappa')\n",
    "ax.grid(True)\n",
    "fig.legend()\n",
    "\n",
    "# Report\n",
    "Y_test_pred_lab = le.inverse_transform(results['Y_pred'])  # back to text labels\n",
    "Y_test_lab = le.inverse_transform(Y_test)  # back to text labels\n",
    "print('\\nClassifier performance')\n",
    "print('Out of sample:\\n', metrics.classification_report(Y_test_lab, Y_test_pred_lab))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise 1**: How do we describe models like this LSTM that have good training and validation performance, but poor test performance? "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise 2**: Try improving the performance of the model. Here are some things to try:\n",
    "- Regularization. Early stopping.\n",
    "- CNN (CNN-LSTM)\n",
    "- Architecture design, optimizer, etc."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.12 ('wearables_workshop')",
   "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.9.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "ba62a1b4a2d59cbe868c83bdd535aa95e1f6d51e6a8dbfe9705911f17ded0148"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}