clinical_mtlr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, TensorDataset, Subset\n",
    "from torch.optim import Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_lightning import seed_everything\n",
    "\n",
    "from sklearn.model_selection import KFold, train_test_split\n",
    "from sklearn.preprocessing import scale\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "from tqdm import tqdm, trange\n",
    "from math import floor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from lifelines.utils import concordance_index\n",
    "from torchmtlr import MTLR, mtlr_neg_log_likelihood, mtlr_cif, mtlr_risk, mtlr_survival\n",
    "from torchmtlr.utils import make_time_bins, encode_survival, reset_parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1129"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seed_everything(1129)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## helper methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "def make_optimizer(opt_cls, model, **kwargs):\n",
    "    \"\"\"Creates a PyTorch optimizer for MTLR training.\"\"\"\n",
    "    params_dict = dict(model.named_parameters())\n",
    "    weights = [v for k, v in params_dict.items() if \"mtlr\" not in k and \"bias\" not in k]\n",
    "    biases = [v for k, v in params_dict.items() if \"bias\" in k]\n",
    "    mtlr_weights = [v for k, v in params_dict.items() if \"mtlr_weight\" in k]\n",
    "    # Don't use weight decay on the biases and MTLR parameters, which have\n",
    "    # their own separate L2 regularization\n",
    "    optimizer = opt_cls([\n",
    "        {\"params\": weights},\n",
    "        {\"params\": biases, \"weight_decay\": 0.},\n",
    "        {\"params\": mtlr_weights, \"weight_decay\": 0.},\n",
    "    ], **kwargs)\n",
    "    return optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "def make_data(path, split=\"training\"):\n",
    "    \"\"\"Load and preprocess the data.\"\"\"\n",
    "    try:\n",
    "        df = pd.read_csv(path)\n",
    "    except:\n",
    "        df = path\n",
    "\n",
    "    clinical_data = (df\n",
    "                     .query(\"split == @split\")\n",
    "                     .set_index(\"Study ID\")\n",
    "                     .drop([\"split\"], axis=1, errors=\"ignore\"))\n",
    "    # if split == \"training\":\n",
    "    clinical_data = clinical_data.rename(columns={\"death\": \"event\", \"survival_time\": \"time\"})\n",
    "    # Convert time to months\n",
    "    clinical_data[\"time\"] *= 12\n",
    "\n",
    "    clinical_data[\"age at dx\"] = scale(clinical_data[\"age at dx\"])\n",
    "    clinical_data[\"Dose\"] = scale(clinical_data[\"Dose\"])\n",
    "\n",
    "    # binarize T stage as T1/2 = 0, T3/4 = 1\n",
    "    clinical_data[\"T Stage\"] = clinical_data[\"T Stage\"].map(\n",
    "        lambda x: \"T1/2\" if x in [\"T1\", \"T1a\", \"T1b\", \"T2\"] else \"T3/4\", na_action=\"ignore\")\n",
    "\n",
    "    # use more fine-grained grouping for N stage\n",
    "    clinical_data[\"N Stage\"] = clinical_data[\"N Stage\"].str.slice(0, 2)\n",
    "\n",
    "    clinical_data[\"Stage\"] = clinical_data[\"Stage\"].map(\n",
    "        lambda x: \"I/II\" if x in [\"I\", \"II\", \"IIA\"] else \"III/IV\", na_action=\"ignore\")\n",
    "\n",
    "    clinical_data[\"ECOG\"] = clinical_data[\"ECOG\"].map(\n",
    "        lambda x: \">0\" if x > 0 else \"0\", na_action=\"ignore\")\n",
    "\n",
    "    clinical_data = pd.get_dummies(clinical_data,\n",
    "                                   columns=[\"Sex\",\n",
    "                                            \"N Stage\",\n",
    "                                            \"Disease Site\"],\n",
    "                                   drop_first=True)\n",
    "    clinical_data = pd.get_dummies(clinical_data,\n",
    "                                   columns=[\"HPV Combined\",\n",
    "                                            \"T Stage\",\n",
    "                                            \"Stage\",\n",
    "                                            \"ECOG\"])\n",
    "\n",
    "    return clinical_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "def multiple_events(row):\n",
    "    event        = row[\"event\"]\n",
    "    cancer_death = row[\"cancer_death\"]\n",
    "\n",
    "    if event==0:\n",
    "        return 0\n",
    "    elif cancer_death==0:\n",
    "        return 1\n",
    "    elif cancer_death==1:\n",
    "        return 2\n",
    "    else:\n",
    "        raise UhOh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "def train_mtlr(x, y, model, time_bins,\n",
    "               num_epochs=1000, lr=.01, weight_decay=0.,\n",
    "               C1=1., batch_size=None,\n",
    "               verbose=True, device=\"cpu\"):\n",
    "    \"\"\"Trains the MTLR model using minibatch gradient descent.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    model : torch.nn.Module\n",
    "        MTLR model to train.\n",
    "    data_train : pd.DataFrame\n",
    "        The training dataset. Must contain a `time` column with the\n",
    "        event time for each sample and an `event` column containing\n",
    "        the event indicator.\n",
    "    num_epochs : int\n",
    "        Number of training epochs.\n",
    "    lr : float\n",
    "        The learning rate.\n",
    "    weight_decay : float\n",
    "        Weight decay strength for all parameters *except* the MTLR\n",
    "        weights. Only used for Deep MTLR training.\n",
    "    C1 : float\n",
    "        L2 regularization (weight decay) strenght for MTLR parameters.\n",
    "    batch_size : int\n",
    "        The batch size.\n",
    "    verbose : bool\n",
    "        Whether to display training progress.\n",
    "    device : str\n",
    "        Device name or ID to use for training.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    torch.nn.Module\n",
    "        The trained model.\n",
    "    \"\"\"\n",
    "    optimizer = make_optimizer(Adam, model, lr=lr, weight_decay=weight_decay)\n",
    "    reset_parameters(model)\n",
    "    print(x.shape, y.shape)\n",
    "    model = model.to(device)\n",
    "    model.train()\n",
    "    train_loader = DataLoader(TensorDataset(x, y), batch_size=batch_size, shuffle=True)\n",
    "    \n",
    "    pbar =  trange(num_epochs, disable=not verbose)\n",
    "    for i in pbar:\n",
    "        for xi, yi in train_loader:\n",
    "            xi, yi = xi.to(device), yi.to(device)\n",
    "            y_pred = model(xi)\n",
    "            loss = mtlr_neg_log_likelihood(y_pred, yi, model, C1, average=True)\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        pbar.set_description(f\"[epoch {i+1: 4}/{num_epochs}]\")\n",
    "        pbar.set_postfix_str(f\"loss = {loss.item():.4f}\")\n",
    "    model.eval()\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "def validate(model, X, time, other, cancer):\n",
    "    \"\"\"Returns inference of model on data X\n",
    "    params\n",
    "    ------\n",
    "    model\n",
    "        PyTorch model being tested\n",
    "    X\n",
    "        data to test on\n",
    "    time, other, cancer\n",
    "        true labels for time, other, and cancer\n",
    "    \"\"\"\n",
    "    pred_prob       = model(X)\n",
    "    two_year_bin    = np.digitize(2, time_bins)\n",
    "    survival_event  = mtlr_survival(pred_prob[:,:29]).detach().numpy()\n",
    "    survival_cancer = mtlr_survival(pred_prob[:,29:]).detach().numpy()\n",
    "    pred_event      = 1 - survival_event[:, two_year_bin]\n",
    "    pred_cancer     = 1 - survival_cancer[:, two_year_bin]\n",
    "    \n",
    "    roc_auc_event   = roc_auc_score(other, pred_event)\n",
    "    roc_auc_cancer  = roc_auc_score(cancer, pred_cancer)\n",
    "\n",
    "    pred_risk = mtlr_risk(pred_prob, 2).detach().numpy()\n",
    "        \n",
    "    ci_event  = concordance_index(time, -pred_risk[:, 0], event_observed=other)\n",
    "    ci_cancer = concordance_index(time, -pred_risk[:, 1], event_observed=cancer)\n",
    "    \n",
    "    return roc_auc_cancer, roc_auc_event, ci_cancer, ci_event"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## load/process data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = \"path/to/data/clinical.csv\"\n",
    "df  = make_data(data_path, split=\"training\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "time_bins = make_time_bins(df[\"time\"], event=df[\"event\"])\n",
    "multi_events = df.apply(lambda x: multiple_events(x), axis=1)\n",
    "\n",
    "y = encode_survival(df[\"time\"], multi_events, time_bins)\n",
    "X = torch.tensor(df.drop([\"time\", \"event\", \"target_binary\", \"cancer_death\"], axis=1).values, dtype=torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_indices = range(len(df))\n",
    "full_targets = df[\"target_binary\"]\n",
    "val_size = floor(.1 / .7 * len(full_indices))\n",
    "\n",
    "train_indices, val_indices = train_test_split(full_indices, test_size=val_size, stratify=full_targets, random_state=1129)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_val = X[train_indices], X[val_indices]\n",
    "y_train, y_val = y[train_indices], y[val_indices]\n",
    "\n",
    "df_val = df.iloc[val_indices]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test = make_data(ddf, split=\"test\")\n",
    "df_test.insert(11, 'N Stage_NX', np.zeros(df_test.shape[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_test = torch.tensor(df_test.drop([\"time\", \"event\", \"target_binary\", \"cancer_death\"], axis=1).values, dtype=torch.float)"
   ]
  },
  {
   "source": [
    "# MTLR training"
   ],
   "cell_type": "markdown",
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linear MTLR fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[epoch    3/350]:   0%|          | 1/350 [00:00<00:37,  9.19it/s, loss = 14.7189]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1545, 24]) torch.Size([1545, 58])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[epoch  350/350]: 100%|██████████| 350/350 [00:12<00:00, 27.19it/s, loss = 1.6799]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.7783771106941839, 0.758510906569159, 0.764453477868112, 0.6833411130128105)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# fit MTLR model \n",
    "mtlr = MTLR(in_features=24, num_time_bins=29, num_events=2)            \n",
    "mtlr = train_mtlr(X_train, y_train, mtlr, time_bins, num_epochs=350, \n",
    "                  lr=1e-3, batch_size=128, verbose=True, device=device, C1=1.)\n",
    "\n",
    "# validation set\n",
    "validate(mtlr, X_val, df_val[\"time\"], df_val[\"event\"], df_val[\"cancer_death\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.7442810457516339,\n",
       " 0.7232062793324715,\n",
       " 0.7470975647112881,\n",
       " 0.6786743839104441)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## Model Performance on test set - REPORTED RESULTS ##\n",
    "validate(mtlr, X_test, df_test[\"time\"], df_test[\"event\"], df_test[\"cancer_death\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neural MTLR\n",
    "* Selected model based on validation set performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[epoch    2/350]:   1%|          | 2/350 [00:00<00:32, 10.82it/s, loss = 31.4964]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1545, 24]) torch.Size([1545, 58])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[epoch  350/350]: 100%|██████████| 350/350 [00:44<00:00,  7.93it/s, loss = 1.8870]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.7621013133208255,\n",
       " 0.7870066826377506,\n",
       " 0.7438572719060524,\n",
       " 0.7153933924588973)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 1 hidden layer\n",
    "\n",
    "mtlr1 = nn.Sequential(\n",
    "    nn.Linear(24, 512),\n",
    "    nn.BatchNorm1d(512),\n",
    "    nn.ReLU(inplace=True),\n",
    "    nn.Dropout(0.2),\n",
    "    \n",
    "    MTLR(512, 29, num_events=2)\n",
    ")\n",
    "\n",
    "mtlr1 = train_mtlr(X_train, y_train, mtlr1, time_bins, num_epochs=350, \n",
    "                   lr=1e-3, batch_size=128, verbose=True, device=device, C1=1.)\n",
    "\n",
    "# validation set\n",
    "validate(mtlr1, X_val, df_val[\"time\"], df_val[\"event\"], df_val[\"cancer_death\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.7676939603001695,\n",
       " 0.7551910468271437,\n",
       " 0.7558891101240619,\n",
       " 0.7256116027862839)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## Model Performance on test set - REPORTED RESULTS ##\n",
    "validate(mtlr1, X_test, df_test[\"time\"], df_test[\"event\"], df_test[\"cancer_death\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[epoch    1/350]:   0%|          | 1/350 [00:00<00:58,  5.94it/s, loss = 42.5312]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1545, 24]) torch.Size([1545, 58])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[epoch  350/350]: 100%|██████████| 350/350 [01:05<00:00,  5.33it/s, loss = 3.3332]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.7000938086303939, 0.7291640398436515, 0.6863143631436315, 0.675509569005757)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 2 hidden layers\n",
    "\n",
    "mtlr2 = nn.Sequential(\n",
    "    nn.Linear(24, 512),\n",
    "    nn.BatchNorm1d(512),\n",
    "    nn.ReLU(inplace=True),\n",
    "    nn.Dropout(0.2),\n",
    "    \n",
    "    nn.Linear(512, 512),\n",
    "    nn.BatchNorm1d(512),\n",
    "    nn.ReLU(inplace=True),\n",
    "    nn.Dropout(0.2),\n",
    "    \n",
    "    MTLR(512, 29, num_events=2)\n",
    ")\n",
    "\n",
    "mtlr2 = train_mtlr(X_train, y_train, mtlr2, time_bins, num_epochs=350, \n",
    "                   lr=1e-3, batch_size=128, verbose=True, device=device, C1=1.)\n",
    "\n",
    "validate(mtlr2, X_val, df_val[\"time\"], df_val[\"event\"], df_val[\"cancer_death\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.7548338779956427,\n",
       " 0.6825163912471323,\n",
       " 0.7146883136774391,\n",
       " 0.6573005350472794)"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## TEST MODEL ##\n",
    "validate(mtlr2, X_test, df_test[\"time\"], df_test[\"event\"], df_test[\"cancer_death\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|          | 0/450 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1545, 24]) torch.Size([1545, 58])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[epoch  450/450]: 100%|██████████| 450/450 [02:02<00:00,  3.67it/s, loss = 1.4988]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.6676829268292683, 0.689509519606607, 0.6225383920505871, 0.6697526061926249)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 3 hidden layers\n",
    "\n",
    "mtlr3 = nn.Sequential(\n",
    "    nn.Linear(24, 512),\n",
    "    nn.BatchNorm1d(512),\n",
    "    nn.ReLU(inplace=True),\n",
    "    nn.Dropout(0.2),\n",
    "    \n",
    "    nn.Linear(512, 512),\n",
    "    nn.BatchNorm1d(512),\n",
    "    nn.ReLU(inplace=True),\n",
    "    nn.Dropout(0.2),\n",
    "    \n",
    "    nn.Linear(512, 512),\n",
    "    nn.BatchNorm1d(512),\n",
    "    nn.ReLU(inplace=True),\n",
    "    nn.Dropout(0.2),\n",
    "    \n",
    "    MTLR(512, 29, num_events=2)\n",
    ")\n",
    "\n",
    "mtlr3 = train_mtlr(X_train, y_train, mtlr3, time_bins, num_epochs=450, \n",
    "                   lr=1e-3, batch_size=128, verbose=True, device=device, C1=1.)\n",
    "\n",
    "validate(mtlr3, X_val, df_val[\"time\"], df_val[\"event\"], df_val[\"cancer_death\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.723160251755023, 0.6591552248192013, 0.6943176596722316, 0.6495384235734877)"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## TEST MODEL ##\n",
    "validate(mtlr3, X_test, df_test[\"time\"], df_test[\"event\"], df_test[\"cancer_death\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "  0%|          | 0/350 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1545, 24]) torch.Size([1545, 58])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[epoch  350/350]: 100%|██████████| 350/350 [01:34<00:00,  3.72it/s, loss = 2.1508]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(0.6674953095684804,\n",
       " 0.7144118017904426,\n",
       " 0.6264227642276423,\n",
       " 0.6774804211399824)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mtlr2 = nn.Sequential(\n",
    "    nn.Linear(24, 512),\n",
    "    nn.BatchNorm1d(512),\n",
    "    nn.ReLU(inplace=True),\n",
    "    #nn.Dropout(0.333),\n",
    "    \n",
    "    nn.Linear(512, 512),\n",
    "    nn.BatchNorm1d(512),\n",
    "    nn.ReLU(inplace=True),\n",
    "    #nn.Dropout(0.333),\n",
    "    \n",
    "    MTLR(512, 29, num_events=2)\n",
    ")\n",
    "\n",
    "mtlr2 = train_mtlr(X_train, y_train, mtlr2, time_bins, num_epochs=350, \n",
    "                   lr=1e-4, batch_size=128, verbose=True, device=device)\n",
    "\n",
    "validate(mtlr2, X_val, df_val[\"time\"], df_val[\"event\"], df_val[\"cancer_death\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.7535781287823772, 0.6954197673199874, 0.7233420125593506, 0.679403483976624)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## TEST MODEL ##\n",
    "validate(mtlr2, X_test, df_test[\"time\"], df_test[\"event\"], df_test[\"cancer_death\"])"
   ]
  }
 ],
 "metadata": {
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}