02_lightning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp lightning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "import logging\n",
    "logging.disable(logging.CRITICAL)\n",
    "import os\n",
    "import random\n",
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "from test_tube import HyperOptArgumentParser\n",
    "from emotion_transformer.dataloader import dataloader\n",
    "from emotion_transformer.model import sentence_embeds_model, context_classifier_model, metrics, f1_score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyTorch Lightning\n",
    "\n",
    "> construction of the PyTorch Lightning module and the hyperparameter search for the SemEval-2019 Task 3 dataset (contextual emotion detection in text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lightning Module\n",
    "\n",
    "Defining the Lightning module is now straightforward, see also the [documentation](https://williamfalcon.github.io/pytorch-lightning/). The default hyperparameter choices were motivated by [this paper](https://arxiv.org/pdf/1905.05583.pdf). \n",
    "\n",
    "Further references for PyTorch Lightning and its usage for Multi-GPU Training/Hyperparameter search can be found in the following blog posts by William Falcon: \n",
    "\n",
    "* [9 Tips For Training Lightning-Fast Neural Networks In Pytorch](https://towardsdatascience.com/9-tips-for-training-lightning-fast-neural-networks-in-pytorch-8e63a502f565)\n",
    "\n",
    "* [Trivial Multi-Node Training With Pytorch-Lightning](https://towardsdatascience.com/trivial-multi-node-training-with-pytorch-lightning-ff75dfb809bd?gi=ec854edcc8eb)\n",
    "\n",
    "* [Converting From Keras To PyTorch Lightning](https://towardsdatascience.com/converting-from-keras-to-pytorch-lightning-be40326d7b7d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class EmotionModel(pl.LightningModule):\n",
    "    \"\"\"\n",
    "    PyTorch Lightning module for the Contextual Emotion Detection in Text Challenge\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, hparams):\n",
    "        \"\"\"\n",
    "        pass in parsed HyperOptArgumentParser to the model\n",
    "        \"\"\"\n",
    "        super(EmotionModel, self).__init__()\n",
    "        self.hparams = hparams\n",
    "        self.emo_dict = {'others': 0, 'sad': 1, 'angry': 2, 'happy': 3}\n",
    "        self.sentence_embeds_model = sentence_embeds_model(dropout = hparams.dropout)\n",
    "        self.context_classifier_model = context_classifier_model(self.sentence_embeds_model.embedding_size,\n",
    "                                                                 hparams.projection_size, \n",
    "                                                                 hparams.n_layers, \n",
    "                                                                 self.emo_dict, \n",
    "                                                                 dropout = hparams.dropout)\n",
    "        \n",
    "\n",
    "    def forward(self, input_ids, attention_mask, labels = None):\n",
    "        \"\"\"\n",
    "        no special modification required for lightning, define as you normally would\n",
    "        \"\"\"\n",
    "        if self.current_epoch < self.hparams.frozen_epochs:\n",
    "            with torch.no_grad():\n",
    "                sentence_embeds = self.sentence_embeds_model(input_ids = input_ids, \n",
    "                                                             attention_mask = attention_mask)\n",
    "        else:\n",
    "            sentence_embeds = self.sentence_embeds_model(input_ids = input_ids, attention_mask = attention_mask)\n",
    "        return self.context_classifier_model(sentence_embeds = sentence_embeds, labels = labels)\n",
    "\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        \"\"\"\n",
    "        Lightning calls this inside the training loop\n",
    "        \"\"\"       \n",
    "        input_ids, attention_mask, labels = batch\n",
    "        loss, _ = self.forward(input_ids = input_ids, attention_mask = attention_mask, labels = labels)\n",
    "        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning\n",
    "        if self.trainer.use_dp or self.trainer.use_ddp2:\n",
    "            loss = loss.unsqueeze(0)\n",
    "\n",
    "        tensorboard_logs = {'train_loss': loss}\n",
    "        return {'loss': loss, 'log': tensorboard_logs}\n",
    "\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        \"\"\"\n",
    "        Lightning calls this inside the validation loop\n",
    "        \"\"\"\n",
    "        input_ids, attention_mask, labels = batch\n",
    "\n",
    "        loss, logits = self.forward(input_ids = input_ids, attention_mask = attention_mask, labels = labels)\n",
    "        scores_dict = metrics(loss, logits, labels)\n",
    "\n",
    "        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning\n",
    "        if self.trainer.use_dp or self.trainer.use_ddp2:\n",
    "            scores = [score.unsqueeze(0) for score in scores_dict.values()]\n",
    "            scores_dict = {key: value for key, value in zip(scores_dict.keys(), scores)}\n",
    "\n",
    "        return scores_dict\n",
    "\n",
    "    def validation_end(self, outputs):\n",
    "        \"\"\"\n",
    "        called at the end of validation to aggregate outputs\n",
    "        :param outputs: list of individual outputs of each validation step\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        \n",
    "        tqdm_dict = {}\n",
    "\n",
    "        for metric_name in outputs[0].keys():\n",
    "            metric_total = 0\n",
    "\n",
    "            for output in outputs:\n",
    "                metric_value = output[metric_name]\n",
    "\n",
    "                if self.trainer.use_dp or self.trainer.use_ddp2:\n",
    "                    if metric_name in ['tp', 'fp', 'fn']:\n",
    "                        metric_value = torch.sum(metric_value)\n",
    "                    else:\n",
    "                        metric_value = torch.mean(metric_value)\n",
    "                    \n",
    "                metric_total += metric_value\n",
    "            if metric_name in ['tp', 'fp', 'fn']:\n",
    "                tqdm_dict[metric_name] = metric_total\n",
    "            else:\n",
    "                tqdm_dict[metric_name] = metric_total / len(outputs)\n",
    "\n",
    "               \n",
    "        prec_rec_f1 = f1_score(tqdm_dict['tp'], tqdm_dict['fp'], tqdm_dict['fn'])\n",
    "        tqdm_dict.update(prec_rec_f1) \n",
    "        result = {'progress_bar': tqdm_dict, 'log': tqdm_dict, 'val_loss': tqdm_dict[\"val_loss\"]}\n",
    "        return result\n",
    "    \n",
    "    def test_step(self, batch, batch_idx):\n",
    "        return self.validation_step(batch, batch_idx)\n",
    "    \n",
    "    def test_end(self, outputs):\n",
    "        return self.validation_end(outputs)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        \"\"\"\n",
    "        returns the optimizer and scheduler\n",
    "        \"\"\"\n",
    "        opt_parameters = self.sentence_embeds_model.layerwise_lr(self.hparams.lr, \n",
    "                                                                 self.hparams.layerwise_decay)\n",
    "        opt_parameters += [{'params': self.context_classifier_model.parameters()}]\n",
    "\n",
    "        optimizer = torch.optim.AdamW(opt_parameters, lr=self.hparams.lr)        \n",
    "        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)\n",
    "        return [optimizer], [scheduler]\n",
    "\n",
    "    \n",
    "    @pl.data_loader\n",
    "    def train_dataloader(self):\n",
    "        return dataloader(self.hparams.train_file, self.hparams.max_seq_len, \n",
    "                          self.hparams.bs, self.emo_dict, use_ddp = self.use_ddp)\n",
    "\n",
    "    \n",
    "    @pl.data_loader\n",
    "    def val_dataloader(self):\n",
    "        return dataloader(self.hparams.val_file, self.hparams.max_seq_len, \n",
    "                          self.hparams.bs, self.emo_dict, use_ddp = self.use_ddp)\n",
    "\n",
    "    \n",
    "    @pl.data_loader\n",
    "    def test_dataloader(self):\n",
    "        return dataloader(self.hparams.test_file, self.hparams.max_seq_len, \n",
    "                          self.hparams.bs, self.emo_dict, use_ddp = self.use_ddp)\n",
    "\n",
    "    \n",
    "    @staticmethod\n",
    "    def add_model_specific_args(parent_parser, root_dir):  \n",
    "        \"\"\"\n",
    "        parameters defined here will be available to the model through self.hparams\n",
    "        \"\"\"\n",
    "        parser = HyperOptArgumentParser(parents=[parent_parser])\n",
    "\n",
    "        parser.opt_list('--bs', default=64, type=int, options=[32, 128, 256], tunable=True,\n",
    "                        help='mini-batch size (default: 256), this is the total batch size of all GPUs'\n",
    "                        'on the current node when using Data Parallel or Distributed Data Parallel')\n",
    "        parser.opt_list('--projection_size', default=256, type=int, options=[32, 128, 512], tunable=True,\n",
    "                       help='sentence embedding size and hidden size for the second transformer')\n",
    "        parser.opt_list('--n_layers', default=1, type=int, options=[2, 4, 6], tunable=True,\n",
    "                       help='number of encoder layers for the second transformer')\n",
    "        parser.opt_list('--frozen_epochs', default=2, type=int, options=[3, 6, 9], tunable=True,\n",
    "                       help='number of epochs the pretrained DistilBert is frozen')\n",
    "        parser.opt_range('--lr', default=2.0e-5, type=float, tunable=True, low=1.0e-5, high=5.0e-4, \n",
    "                         nb_samples=5, help='initial learning rate')\n",
    "        parser.opt_list('--layerwise_decay', default=0.95, type=float, options=[0.3, 0.6, 0.8], tunable=True,\n",
    "                       help='layerwise decay factor for the learning rate of the pretrained DistilBert')\n",
    "        parser.opt_list('--max_seq_len', default=32, type=int, options=[16, 64], tunable=False,\n",
    "                       help='maximal number of input tokens for the DistilBert model')\n",
    "        parser.opt_list('--dropout', default=0.1, type=float, options=[0.1, 0.2], tunable=False)\n",
    "        parser.add_argument('--train_file', default=os.path.join(root_dir, 'data/clean_train.txt'), type=str)\n",
    "        parser.add_argument('--val_file', default=os.path.join(root_dir, 'data/clean_val.txt'), type=str)\n",
    "        parser.add_argument('--test_file', default=os.path.join(root_dir, 'data/clean_test.txt'), type=str)\n",
    "        parser.add_argument('--epochs', default=10, type=int, metavar='N',\n",
    "                            help='number of total epochs to run')\n",
    "        parser.add_argument('--seed', type=int, default=None,\n",
    "                            help='seed for initializing training')\n",
    "        \n",
    "        return parser"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperparameter Search Argument Parser\n",
    "\n",
    "Next we define the HyperOptArgumentParser including distributed training (see also the [documentation](https://williamfalcon.github.io/pytorch-lightning/Trainer/Distributed%20training/\n",
    ")) and debugging functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def get_args(model):\n",
    "    \"\"\"\n",
    "    returns the HyperOptArgumentParser\n",
    "    \"\"\"\n",
    "    parent_parser = HyperOptArgumentParser(strategy='random_search', add_help = False)\n",
    "\n",
    "    root_dir = os.getcwd()  \n",
    "    parent_parser.add_argument('--mode', type=str, default='default', \n",
    "                               choices=('default', 'test', 'hparams_search'),\n",
    "                               help='supports default for train/test/val and hparams_search for a hyperparameter search')\n",
    "    parent_parser.add_argument('--save-path', metavar='DIR', default=os.path.join(root_dir, 'logs'), type=str,\n",
    "                               help='path to save output')\n",
    "    parent_parser.add_argument('--gpus', type=str, default=None, help='which gpus')\n",
    "    parent_parser.add_argument('--distributed-backend', type=str, default=None, choices=('dp', 'ddp', 'ddp2'),\n",
    "                               help='supports three options dp, ddp, ddp2')\n",
    "    parent_parser.add_argument('--use_16bit', dest='use_16bit', action='store_true',\n",
    "                               help='if true uses 16 bit precision')\n",
    "\n",
    "    # debugging\n",
    "    parent_parser.add_argument('--fast_dev_run', dest='fast_dev_run', action='store_true',\n",
    "                               help='debugging a full train/val/test loop')\n",
    "    parent_parser.add_argument('--track_grad_norm', dest='track_grad_norm', action='store_true',\n",
    "                               help='inspect gradient norms')\n",
    "\n",
    "    parser = model.add_model_specific_args(parent_parser, root_dir) \n",
    "    return parser"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us take a look at the different attributes of `hparams`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'mode': 'default',\n",
       " 'save_path': '/home/julius/Documents/nbdev_venv/emotion_transformer/logs',\n",
       " 'gpus': None,\n",
       " 'distributed_backend': None,\n",
       " 'use_16bit': False,\n",
       " 'fast_dev_run': False,\n",
       " 'track_grad_norm': False,\n",
       " 'bs': 64,\n",
       " 'projection_size': 256,\n",
       " 'n_layers': 1,\n",
       " 'frozen_epochs': 2,\n",
       " 'lr': 2e-05,\n",
       " 'layerwise_decay': 0.95,\n",
       " 'max_seq_len': 32,\n",
       " 'dropout': 0.1,\n",
       " 'train_file': '/home/julius/Documents/nbdev_venv/emotion_transformer/data/clean_train.txt',\n",
       " 'val_file': '/home/julius/Documents/nbdev_venv/emotion_transformer/data/clean_val.txt',\n",
       " 'test_file': '/home/julius/Documents/nbdev_venv/emotion_transformer/data/clean_test.txt',\n",
       " 'epochs': 10,\n",
       " 'seed': None,\n",
       " 'hpc_exp_number': None,\n",
       " 'trials': <bound method HyperOptArgumentParser.opt_trials of HyperOptArgumentParser(prog='ipykernel_launcher.py', usage=None, description=None, formatter_class=<class 'argparse.HelpFormatter'>, conflict_handler='error', add_help=True)>,\n",
       " 'optimize_parallel': <bound method HyperOptArgumentParser.optimize_parallel of HyperOptArgumentParser(prog='ipykernel_launcher.py', usage=None, description=None, formatter_class=<class 'argparse.HelpFormatter'>, conflict_handler='error', add_help=True)>,\n",
       " 'optimize_parallel_gpu': <bound method HyperOptArgumentParser.optimize_parallel_gpu of HyperOptArgumentParser(prog='ipykernel_launcher.py', usage=None, description=None, formatter_class=<class 'argparse.HelpFormatter'>, conflict_handler='error', add_help=True)>,\n",
       " 'optimize_parallel_cpu': <bound method HyperOptArgumentParser.optimize_parallel_cpu of HyperOptArgumentParser(prog='ipykernel_launcher.py', usage=None, description=None, formatter_class=<class 'argparse.HelpFormatter'>, conflict_handler='error', add_help=True)>,\n",
       " 'generate_trials': <bound method HyperOptArgumentParser.generate_trials of HyperOptArgumentParser(prog='ipykernel_launcher.py', usage=None, description=None, formatter_class=<class 'argparse.HelpFormatter'>, conflict_handler='error', add_help=True)>,\n",
       " 'optimize_trials_parallel_gpu': <bound method HyperOptArgumentParser.optimize_trials_parallel_gpu of HyperOptArgumentParser(prog='ipykernel_launcher.py', usage=None, description=None, formatter_class=<class 'argparse.HelpFormatter'>, conflict_handler='error', add_help=True)>}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hparams = get_args(EmotionModel)\n",
    "hparams = hparams.parse_args(args=[])\n",
    "vars(hparams)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trainer\n",
    "Next we define a function calling the Lightning trainer using the setting specified in `hparams`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def main(hparams, gpus = None):\n",
    "    \"\"\"\n",
    "    Trains the Lightning model as specified in `hparams`\n",
    "    \"\"\"\n",
    "    model = EmotionModel(hparams)\n",
    "    \n",
    "    if hparams.seed is not None:\n",
    "        random.seed(hparams.seed)\n",
    "        torch.manual_seed(hparams.seed)\n",
    "        torch.backends.cudnn.deterministic = True\n",
    "        \n",
    "    early_stop_callback = pl.callbacks.EarlyStopping(monitor='val_loss', min_delta=0.00, patience=5, \n",
    "                                        verbose=False, mode='min')\n",
    "    \n",
    "    \n",
    "    trainer = pl.Trainer(default_save_path=hparams.save_path,\n",
    "                        gpus=len(gpus.split(\",\")) if gpus else hparams.gpus,\n",
    "                        distributed_backend=hparams.distributed_backend,\n",
    "                        use_amp=hparams.use_16bit,\n",
    "                        early_stop_callback=early_stop_callback,\n",
    "                        max_nb_epochs=hparams.epochs,\n",
    "                        fast_dev_run=hparams.fast_dev_run,\n",
    "                        track_grad_norm=(2 if hparams.track_grad_norm else -1))\n",
    "    trainer.fit(model)\n",
    "    \n",
    "    if hparams.mode == 'test':\n",
    "        trainer.test()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us check the model by running a quick development run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1:  50%|█████     | 1/2 [00:03<00:03,  3.96s/batch, batch_nb=0, loss=2.086, v_nb=13]\n",
      "Validating:   0%|          | 0/1 [00:00<?, ?batch/s]\u001b[A\n",
      "Epoch 1: 100%|██████████| 2/2 [00:07<00:00,  3.96s/batch, batch_nb=0, f1_score=0.13, fn=7, fp=33, loss=2.086, precision=0.0833, recall=0.3, tp=3, v_nb=13, val_acc=0.406, val_loss=2.03]\n",
      "Epoch 1: 100%|██████████| 2/2 [00:08<00:00,  4.22s/batch, batch_nb=0, f1_score=0.13, fn=7, fp=33, loss=2.086, precision=0.0833, recall=0.3, tp=3, v_nb=13, val_acc=0.406, val_loss=2.03]\n"
     ]
    }
   ],
   "source": [
    "hparams.fast_dev_run = True\n",
    "main(hparams)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also create a python file for automatic hyperparameter optimization across different GPUs or CPUs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting main.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile main.py\n",
    "\n",
    "from emotion_transformer.lightning import EmotionModel, get_args, main\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    hparams = get_args(EmotionModel)\n",
    "    hparams = hparams.parse_args()\n",
    "\n",
    "    if hparams.mode in ['test','default']:\n",
    "        main(hparams)\n",
    "    elif hparams.mode == 'hparams_search':\n",
    "        if hparams.gpus:\n",
    "            hparams.optimize_parallel_gpu(main, max_nb_trials=20, \n",
    "                                          gpu_ids = [gpus for gpus in hparams.gpus.split(' ')])\n",
    "        else:\n",
    "            hparams.optimize_parallel_cpu(main, nb_trials=20, nb_workers=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background Information\n",
    "\n",
    "For the interested reader we provide some background information on the (distributed) training loop:\n",
    "\n",
    "* one epoch consists of m = ceil(30160/batchsize) batches for the training and additional n = ceil(2755/batchsize) \n",
    "batches for the validation.\n",
    "\n",
    "**dp case:** \n",
    "\n",
    "* the batchsize will be split and each gpu receives (up to rounding) a batch of size batchsize/num_gpus\n",
    "\n",
    "* in the validation steps each gpu computes its own scores for each of the n batches (of size batchsize/num_gpus), i.e. each gpu calls the `validation_step` method\n",
    "\n",
    "* the `output` which is passed to the `validation_end` method consists of list of dictionaries (containing the concatenated scores from the different gpus), i.e.\n",
    "\n",
    "`output = [ {first_metric: [first_gpu_batch_1,...,last_gpu_batch_1],...,\n",
    "             last_metric:  [first_gpu_batch_1,...,last_gpu_batch_1]},..., \n",
    "            {first_metric: [first_gpu_batch_n,...,last_gpu_batch_n],...,\n",
    "             last_metric:  [first_gpu_batch_n,...,last_gpu_batch_n]} ]`\n",
    "\n",
    "\n",
    "**ddp case:** (does not work from jupyter notebooks)\n",
    "\n",
    "* the gpus receive (disjoint) samples of size batchsize and train on own processes but communicate and average their gradients (thus the resulting models on each gpu have the same weights)\n",
    "\n",
    "* each gpu computes its own validation_end method and its own list of dictionaries \n",
    "\n",
    "`output_first_gpu = [ {first_metric: batch_1,...,last_metric: batch_1},..., \n",
    "                      {first_metric: batch_n,...,last_metric: batch_n} ]`\n",
    "                      \n",
    "`output_last_gpu = [ {first_metric: batch_1,...,last_metric: batch_1},..., \n",
    "                      {first_metric: batch_n,...,last_metric: batch_n} ]`\n",
    "\n",
    "\n",
    "**ddp case:** (does not work from jupyter notebooks)\n",
    "\n",
    "*  on each node we have the dp case but the nodes communicate analogous to the ddp case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted 00_dataloader.ipynb.\n",
      "Converted 01_model.ipynb.\n",
      "Converted 02_lightning.ipynb.\n",
      "Converted index.ipynb.\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "from nbdev.export import notebook2script\n",
    "notebook2script()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}