multimodal_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multimodal deep learning \n",
    "\n",
    "\n",
    "It employes the spectral and tabular deep learning models as modules, which will be trained together to inform the grain yield prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# Import libraries\n",
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "from fastai.vision.all import *\n",
    "import fastai\n",
    "from fastai.tabular.all import *\n",
    "from fastai.data.load import _FakeLoader, _loaders\n",
    "import torch\n",
    "from ipywidgets import IntProgress\n",
    "from glob import glob\n",
    "\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import mean_squared_error, r2_score\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "\n",
    "# Custom functions\n",
    "from msi_utils import *\n",
    "from fold_utils import * \n",
    "from multimodal_utisl import *\n",
    "from multimodal_model import *\n",
    "\n",
    "# use gpu by default if available\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fastai and fastcore versions have to be above \n",
    "\n",
    "fastai : 2.1.2 (at the moment 2.3.0)\n",
    "fastcore: 1.3.1 (1.3.19)\n",
    "torch : 1.7.0 ('1.8.1+cu102')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Reference tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "hidden": true
   },
   "outputs": [],
   "source": [
    "path = Path('/data/g2f_data/input_data/')\n",
    "df_test = pd.read_csv('/data/fielddata/df_test.csv')\n",
    "df_train_val = pd.read_csv('/data/fielddata/df_train_val.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Use random splitter function from fastai\n",
    "splitter = RandomSplitter(seed=42)\n",
    "splits = splitter(range_of(df_train_val))\n",
    "splits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Tab Dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "procs = [Categorify, Normalize, FillMissing]\n",
    "cat_names = ['Parental 1', 'Parental 2', 'Planting', 'Stock', 'Fertilizer']\n",
    "cont_names =['Days_after_sowing']\n",
    "\n",
    "to = TabularPandas(df_train_val,\n",
    "                   procs,\n",
    "                   cat_names,\n",
    "                   cont_names=cont_names,\n",
    "                   y_names='Yield',\n",
    "                   y_block=RegressionBlock(),\n",
    "                   splits=splits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "tab_dl = to.dataloaders(bs=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tab_dl.show_batch()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Spectral Dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "dblock = DataBlock(blocks    = (MSITensorBlock, RegressionBlock),\n",
    "                       get_items = get_npy,\n",
    "                       get_y     = get_y,\n",
    "                       splitter  = splitter)\n",
    "\n",
    "msi_dls = dblock.dataloaders(df_train_val, bs=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "msi_dls.show_batch(channels=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "msi_dls.show_batch(channels=9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Mixed Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Check that the tabular dataset is aligned with the spectral dataset\n",
    "mixed_dl = MixedDL(tab_dl[0], msi_dls[0])\n",
    "# These should show the same ids\n",
    "msi_dls[0].get_idxs()[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Now mix the tabular and spectral datasets to create the multimodal input\n",
    "train_mixed_dl = MixedDL(tab_dl[0], msi_dls[0])\n",
    "valid_mixed_dl = MixedDL(tab_dl[1], msi_dls[1])\n",
    "mixed_dls = DataLoaders(train_mixed_dl, valid_mixed_dl).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "mixed_dls.show_batch()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fusion at feature level - Training modules from scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Kfold "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     27,
     37,
     50,
     60
    ],
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# KFOLD VALIDATION\n",
    "kfold_preds = pd.DataFrame(columns=['predictions', 'target_yield'])\n",
    "\n",
    "# tab variables\n",
    "procs = [Categorify, Normalize, FillMissing]\n",
    "cat_names = ['Parental 1', 'Parental 2', 'Planting', 'Stock', 'Fertilizer']\n",
    "cont_names =['Days_after_sowing']\n",
    "\n",
    "split_list = kfold_splitter(df=df_train_val)\n",
    "\n",
    "# Mixed model variables\n",
    "# Set weights for each loss\n",
    "tab_w, vis_w, tv_w = 0.1, 0.55, 0.35\n",
    "\n",
    "# Initialise Loss\n",
    "gb_loss = myGradientBlending(tab_weight=tab_w, visual_weight=vis_w, tab_vis_weight=tv_w, \n",
    "                             loss_scale=1.0)\n",
    "\n",
    "# METRICS\n",
    "metrics = [t_rmse, v_rmse, tv_rmse, weighted_RMSEp]\n",
    "csvlogger = CSVLogger('/data/results/multimodal_5fold_metrics.csv', append=True)\n",
    "cbs = [csvlogger]    \n",
    "\n",
    "for i in range(5):\n",
    "    getter = get_fold(split_list, fold=i)\n",
    "    splits = getter(range_of(df_train_val))\n",
    "\n",
    "    to = TabularPandas(df_train_val,\n",
    "                   procs,\n",
    "                   cat_names,\n",
    "                   cont_names=cont_names,\n",
    "                   y_names='Yield',\n",
    "                   y_block=RegressionBlock(),\n",
    "                   splits=splits)\n",
    "    tab_dl = to.dataloaders(bs=8)\n",
    "    \n",
    "    # Call MSI dataloader\n",
    "    msi_fold = DataBlock(blocks = (MSITensorBlock, RegressionBlock),\n",
    "                          get_items = get_npy,\n",
    "                          get_y = get_y,\n",
    "                          splitter = getter)\n",
    "    msi_dl = msi_fold.dataloaders(df_train_val, bs=8)\n",
    "\n",
    "    # Make mixed dls\n",
    "    train_mixed_dl = MixedDL(tab_dl[0], msi_dl[0])\n",
    "    valid_mixed_dl = MixedDL(tab_dl[1], msi_dl[1])\n",
    "    mixed_dls = DataLoaders(train_mixed_dl, valid_mixed_dl).cuda()\n",
    "    \n",
    "    # Modules\n",
    "    config = tabular_config(ps=0.5, embed_p=0.5)\n",
    "    learn_tab = tabular_learner(tab_dl,\n",
    "                            config=config,\n",
    "                            layers=[200,100],\n",
    "                            metrics=[rmse, R2Score()],\n",
    "                            opt_func=ranger,\n",
    "                            y_range=[0,20],\n",
    "                            wd=0.3)\n",
    "    learn_tab.fit_one_cycle(1, 1e-3)\n",
    "    \n",
    "    model_msi = xresnet18(n_out=1, c_in=13, pretrained=False, sa=True, p=0.5, ndim=2)\n",
    "    learn_msi = Learner(msi_dl, \n",
    "                model_msi,\n",
    "                opt_func=Adam, \n",
    "                loss_func=root_mean_squared_error,\n",
    "                metrics=[rmse, R2Score()])\n",
    "    learn_msi.fit_one_cycle(1, 1e-3)\n",
    "    \n",
    "    multi_model = TabVis(learn_tab.model, learn_msi.model)\n",
    "    multi_learn = Learner(mixed_dls, multi_model, gb_loss, cbs=cbs, metrics=metrics)\n",
    "    \n",
    "    # Disable Fastai progress bar\n",
    "    with multi_learn.no_bar()and multi_learn.no_logging():\n",
    "        multi_learn.fit_one_cycle(35, lr_max=1e-3)\n",
    "        \n",
    "    preds,targs = multi_learn.get_preds(dl=valid_mixed_dl)\n",
    "    pred_mixed_df = pd.DataFrame()\n",
    "    tab_pred = preds[0].flatten()\n",
    "    vis_pred = preds[1].flatten()\n",
    "    mixed_pred = preds[2].flatten()\n",
    "\n",
    "    pred_mixed_df['tab_pred'] = tab_pred\n",
    "    pred_mixed_df['msi_pred'] = vis_pred\n",
    "    pred_mixed_df['mixed_pred'] = mixed_pred\n",
    "    pred_mixed_df['target_yield'] = targs\n",
    "    kfold_preds = kfold_preds.append(pred_mixed_df)\n",
    "\n",
    "kfold_preds.to_csv('/data/results/multimodal_5fold_prediction_for metrics.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model to predict the holdout dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Train a model for test prediction\n",
    "config = tabular_config(ps=0.5, embed_p=0.5)\n",
    "learn_tab = tabular_learner(tab_dl,\n",
    "                            config=config,\n",
    "                            layers=[200,100],\n",
    "                            metrics=[rmse, R2Score()],\n",
    "                            opt_func=ranger,\n",
    "                            y_range=[0,20],\n",
    "                            wd=0.3)\n",
    "learn_tab.fit_one_cycle(1, 1e-3)\n",
    "\n",
    "model_msi = xresnet18(n_out=1, c_in=13, pretrained=False, sa=True, p=0.5, ndim=2)\n",
    "learn_msi = Learner(msi_dls, \n",
    "                model_msi,\n",
    "                opt_func=Adam, \n",
    "                loss_func=root_mean_squared_error,\n",
    "                metrics=[rmse, R2Score()])\n",
    "learn_msi.fit_one_cycle(1, 1e-3)\n",
    "\n",
    "multi_model = TabVis(learn_tab.model, learn_msi.model)\n",
    "tab_w, vis_w, tv_w = 0.05, 0.5, 0.35 # Set weights for each loss\n",
    "\n",
    "# Initialise Loss\n",
    "gb_loss = myGradientBlending(tab_weight=tab_w, visual_weight=vis_w, tab_vis_weight=tv_w, \n",
    "                             loss_scale=1.0)\n",
    "# Define metrics weights\n",
    "metrics = [t_rmse, v_rmse, tv_rmse, weighted_RMSEp]\n",
    "early_stopping = EarlyStoppingCallback(monitor='valid_loss', patience=3, min_delta=0.01)\n",
    "\n",
    "multi_learn = Learner(mixed_dls, multi_model, gb_loss, metrics=metrics, cbs=early_stopping)\n",
    "multi_learn.fit_one_cycle(60, lr_max=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "multi_learn.recorder.plot_loss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Saved the updated model\n",
    "multi_learn.save('/data/model_weights/multimodal_colearning')\n",
    "\n",
    "# Use this to load the model\n",
    "# multi_learn.load('/data/model_weights/multimodal_colearning')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Holdout dataset results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_msi_dls = dblock.dataloaders(df_test, shuffle=False)\n",
    "learn_msi.dls.loaders.append(msi_dls.test_dl(test_msi_dls.items, with_labels=True, shuffle=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1st half -\n",
    "# Find the order of samples in the MSI test DL\n",
    "\n",
    "fnames_MSIorder =[]\n",
    "for fname in test_msi_dls.items:\n",
    "    fname = str(fname)\n",
    "    fname = fname.split(sep='/')[-1]\n",
    "    fname = fname.replace('.npy', '')\n",
    "    fnames_MSIorder.append(fname)\n",
    "    \n",
    "fnames_MSIorder\n",
    "\n",
    "# Reorder the df_test to reflect this order\n",
    "df_test1 = df_test.set_index('Barcode').reindex(fnames_MSIorder)\n",
    "\n",
    "learn_tab.dls.loaders.append(tab_dl.test_dl(df_test1, with_labels=True, shuffle=False))\n",
    "test_mixed_dl = MixedDL(learn_tab.dls[2], learn_msi.dls[2])\n",
    "test_mixed_dl.show_batch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "preds,targs = multi_learn.get_preds(dl=test_mixed_dl)\n",
    "tab_pred = preds[0].flatten()\n",
    "vis_pred = preds[1].flatten()\n",
    "mixed_pred = preds[2].flatten()\n",
    "\n",
    "mixed_results = df_test1.copy()\n",
    "mixed_results['tab_pred'] = tab_pred\n",
    "mixed_results['msi_pred'] = vis_pred\n",
    "mixed_results['mixed_pred'] = mixed_pred\n",
    "\n",
    "len(mixed_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2nd half -\n",
    "# Find the order of samples in the MSI test DL\n",
    "\n",
    "fnames_MSIorder =[]\n",
    "for fname in test_msi_dls[1].items:\n",
    "    fname = str(fname)\n",
    "    fname = fname.split(sep='/')[-1]\n",
    "    fname = fname.replace('.npy', '')\n",
    "    fnames_MSIorder.append(fname)\n",
    "    \n",
    "# fnames_MSIorder\n",
    "\n",
    "# Reorder the df_test to reflect this order\n",
    "df_test2 = df_test.set_index('Barcode').reindex(fnames_MSIorder)\n",
    "learn_tab.dls.loaders.append(tab_dl.test_dl(df_test2, with_labels=True, shuffle=False))\n",
    "learn_msi.dls.loaders.append(msi_dls.test_dl(test_msi_dls[1].items, with_labels=True, shuffle=False))\n",
    "test_mixed_dl = MixedDL(learn_tab.dls[4], learn_msi.dls[3])\n",
    "test_mixed_dl.show_batch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preds,targs = multi_learn.get_preds(dl=test_mixed_dl)\n",
    "tab_pred = preds[0].flatten()\n",
    "vis_pred = preds[1].flatten()\n",
    "mixed_pred = preds[2].flatten()\n",
    "\n",
    "mixed_results2 = df_test2.copy()\n",
    "mixed_results2['tab_pred'] = tab_pred\n",
    "mixed_results2['msi_pred'] = vis_pred\n",
    "mixed_results2['mixed_pred'] = mixed_pred\n",
    "\n",
    "len(mixed_results2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ff_GB_results = mixed_results.append(mixed_results2)\n",
    "ff_GB_results = mixed_results\n",
    "ff_GB_results.to_csv('/data/results/multimodal_colearning_prediction_on_holdout_dataset.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fusion at feature level - Pretrained modules "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Kfold "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     26,
     36,
     49,
     60
    ],
    "hidden": true
   },
   "outputs": [],
   "source": [
    "kfold_preds = pd.DataFrame(columns=['predictions', 'target_yield'])\n",
    "\n",
    "# tab variables\n",
    "procs = [Categorify, Normalize, FillMissing]\n",
    "cat_names = ['Parental 1', 'Parental 2', 'Planting', 'Stock', 'Fertilizer']\n",
    "cont_names =['Days_after_sowing']\n",
    "\n",
    "split_list = kfold_splitter(df=df_train_val)\n",
    "\n",
    "# Mixed model variables\n",
    "# Set weights for each loss\n",
    "tab_w, vis_w, tv_w = 0.1, 0.55, 0.35\n",
    "\n",
    "# Initialise Loss\n",
    "gb_loss = myGradientBlending(tab_weight=tab_w, visual_weight=vis_w, tab_vis_weight=tv_w, \n",
    "                             loss_scale=1.0)\n",
    "\n",
    "# METRICS\n",
    "metrics = [t_rmse, v_rmse, tv_rmse, weighted_RMSEp]\n",
    "csvlogger = CSVLogger('/data/results/multimodal_pretrained_5fold_metrics.csv', append=True)\n",
    "cbs = [csvlogger]\n",
    "    \n",
    "for i in range(5):\n",
    "    getter = get_fold(split_list, fold=i)\n",
    "    splits = getter(range_of(df_train_val))\n",
    "\n",
    "    to = TabularPandas(df_train_val,\n",
    "                   procs,\n",
    "                   cat_names,\n",
    "                   cont_names=cont_names,\n",
    "                   y_names='Yield',\n",
    "                   y_block=RegressionBlock(),\n",
    "                   splits=splits)\n",
    "    tab_dl = to.dataloaders(bs=8)\n",
    "    \n",
    "    # Call MSI dataloader\n",
    "    msi_fold = DataBlock(blocks = (MSITensorBlock, RegressionBlock),\n",
    "                          get_items = get_npy,\n",
    "                          get_y = get_y,\n",
    "                          splitter = getter)\n",
    "    msi_dl = msi_fold.dataloaders(df_train_val, bs=8)\n",
    "\n",
    "    # Make mixed dls\n",
    "    train_mixed_dl = MixedDL(tab_dl[0], msi_dl[0])\n",
    "    valid_mixed_dl = MixedDL(tab_dl[1], msi_dl[1])\n",
    "    mixed_dls = DataLoaders(train_mixed_dl, valid_mixed_dl).cuda()\n",
    "    \n",
    "    # Modules\n",
    "    config = tabular_config(ps=0.5, embed_p=0.5)\n",
    "    learn_tab = tabular_learner(tab_dl,\n",
    "                            config=config,\n",
    "                            layers=[200,100],\n",
    "                            metrics=[rmse, R2Score()],\n",
    "                            opt_func=ranger,\n",
    "                            y_range=[0,20],\n",
    "                            wd=0.3)\n",
    "    learn_tab.load('/data/model_weights/DNN_model')\n",
    "    learn_tab.fit_one_cycle(1, 1e-3)\n",
    "    \n",
    "    model_msi = xresnet18(n_out=1, c_in=13, pretrained=False, sa=True, p=0.5, ndim=2)\n",
    "    learn_msi = Learner(msi_dl, \n",
    "                model_msi,\n",
    "                opt_func=Adam, \n",
    "                loss_func=root_mean_squared_error,\n",
    "                metrics=[rmse, R2Score()])\n",
    "    learn_msi.load('/data/model_weights/xresnet18_model')\n",
    "    learn_msi.fit_one_cycle(1, 1e-3)\n",
    "    \n",
    "    multi_model = TabVis(learn_tab.model, learn_msi.model)\n",
    "    multi_learn = Learner(mixed_dls, multi_model, gb_loss, cbs=cbs, metrics=metrics)\n",
    "    \n",
    "        # Disable Fastai progress bar\n",
    "    with multi_learn.no_bar()and multi_learn.no_logging():\n",
    "        multi_learn.fit_one_cycle(35, lr_max=1e-3)\n",
    "        \n",
    "    preds,targs = multi_learn.get_preds(dl=valid_mixed_dl)\n",
    "    \n",
    "    pred_mixed_df = pd.DataFrame()\n",
    "    tab_pred = preds[0].flatten()\n",
    "    vis_pred = preds[1].flatten()\n",
    "    mixed_pred = preds[2].flatten()\n",
    "\n",
    "    pred_mixed_df['tab_pred'] = tab_pred\n",
    "    pred_mixed_df['msi_pred'] = vis_pred\n",
    "    pred_mixed_df['mixed_pred'] = mixed_pred\n",
    "    \n",
    "    pred_mixed_df['target_yield'] = targs\n",
    "    \n",
    "    kfold_preds = kfold_preds.append(pred_mixed_df)\n",
    "\n",
    "kfold_preds.to_csv('/data/results/multimodal_pretrained_5fold_prediction_for_metrics.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model to predict the holdout dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0,
     2
    ]
   },
   "outputs": [],
   "source": [
    "# TAB pretrained version\n",
    "config = tabular_config(ps=0.5, embed_p=0.5)\n",
    "preTAB_learner = tabular_learner(tab_dl,\n",
    "                            config=config,\n",
    "                            layers=[200,100],\n",
    "                            metrics=[rmse, R2Score()],\n",
    "                            opt_func=ranger,\n",
    "                            y_range=[0,20],\n",
    "                            wd=0.3)\n",
    "\n",
    "preTAB_learner.load('/data/model_weights/DNN_model')\n",
    "# test that the model is working\n",
    "preTAB_learner.fit_one_cycle(1, 1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# Spectral pretrained version\n",
    "model_msi = xresnet18(n_out=1, c_in=13, pretrained=False, sa=True, p=0.5, ndim=2)\n",
    "preVIS_learner = Learner(msi_dls, \n",
    "                model_msi,\n",
    "                opt_func=Adam, \n",
    "                loss_func=root_mean_squared_error,\n",
    "                metrics=[rmse, R2Score()])\n",
    "\n",
    "preVIS_learner.load('/data/model_weights/xresnet18_model')\n",
    "# test that the model is working\n",
    "preVIS_learner.fit_one_cycle(1, 1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# Train the multimodal model\n",
    "multi_model = TabVis(preTAB_learner.model, preVIS_learner.model)\n",
    "\n",
    "# Set weights for each loss\n",
    "tab_w, vis_w, tv_w = 0.1, 0.55, 0.35\n",
    "\n",
    "# Initialise Loss\n",
    "gb_loss = myGradientBlending(tab_weight=tab_w, visual_weight=vis_w, tab_vis_weight=tv_w, \n",
    "                             loss_scale=1.0)\n",
    "\n",
    "# Define metrics weights\n",
    "metrics = [t_rmse, v_rmse, tv_rmse, weighted_RMSEp]\n",
    "multi_learn = Learner(mixed_dls, multi_model, gb_loss, metrics=metrics)\n",
    "\n",
    "multi_learn.fit_one_cycle(60, lr_max=1e-3)\n",
    "# multi_learn.remove_my_hooks()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Saved the updated model\n",
    "multi_learn.save('/data/model_weights/multimodal_pretrained')\n",
    "# Load the updated model\n",
    "# multi_learn.load('/data/model_weights/multimodal_pretrained')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Holdout dataset results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_msi_dls = dblock.dataloaders(df_test, shuffle=False)\n",
    "preVIS_learner.dls.loaders.append(msi_dls.test_dl(test_msi_dls.items, with_labels=True, shuffle=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1st half -\n",
    "# Find the order of samples in the MSI test DL\n",
    "\n",
    "fnames_MSIorder =[]\n",
    "for fname in test_msi_dls.items:\n",
    "    fname = str(fname)\n",
    "    fname = fname.split(sep='/')[-1]\n",
    "    fname = fname.replace('.npy', '')\n",
    "    fnames_MSIorder.append(fname)\n",
    "    \n",
    "fnames_MSIorder\n",
    "\n",
    "# Reorder the df_test to reflect this order\n",
    "df_test1 = df_test.set_index('Barcode').reindex(fnames_MSIorder)\n",
    "preTAB_learner.dls.loaders.append(tab_dl.test_dl(df_test1, with_labels=True, shuffle=False))\n",
    "test_mixed_dl = MixedDL(preTAB_learner.dls[2], preVIS_learner.dls[2])\n",
    "test_mixed_dl.show_batch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preds,targs = multi_learn.get_preds(dl=test_mixed_dl)\n",
    "tab_pred = preds[0].flatten()\n",
    "vis_pred = preds[1].flatten()\n",
    "mixed_pred = preds[2].flatten()\n",
    "\n",
    "mixed_results = df_test1.copy()\n",
    "mixed_results['tab_pred'] = tab_pred\n",
    "mixed_results['msi_pred'] = vis_pred\n",
    "mixed_results['mixed_pred'] = mixed_pred\n",
    "\n",
    "mixed_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2nd half -\n",
    "# Find the order of samples in the MSI test DL\n",
    "\n",
    "fnames_MSIorder =[]\n",
    "for fname in test_msi_dls[1].items:\n",
    "    fname = str(fname)\n",
    "    fname = fname.split(sep='/')[-1]\n",
    "    fname = fname.replace('.npy', '')\n",
    "    fnames_MSIorder.append(fname)\n",
    "    \n",
    "# fnames_MSIorder\n",
    "\n",
    "# Reorder the df_test to reflect this order\n",
    "df_test2 = df_test.set_index('Barcode').reindex(fnames_MSIorder)\n",
    "preTAB_learner.dls.loaders.append(tab_dl.test_dl(df_test2, with_labels=True, shuffle=False))\n",
    "preVIS_learner.dls.loaders.append(msi_dls.test_dl(test_msi_dls[1].items, with_labels=True, shuffle=False))\n",
    "test_mixed_dl = MixedDL(preTAB_learner.dls[3], preVIS_learner.dls[3])\n",
    "test_mixed_dl.show_batch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preds,targs = multi_learn.get_preds(dl=test_mixed_dl)\n",
    "tab_pred = preds[0].flatten()\n",
    "vis_pred = preds[1].flatten()\n",
    "mixed_pred = preds[2].flatten()\n",
    "\n",
    "mixed_results2 = df_test2.copy()\n",
    "mixed_results2['tab_pred'] = tab_pred\n",
    "mixed_results2['msi_pred'] = vis_pred\n",
    "mixed_results2['mixed_pred'] = mixed_pred\n",
    "\n",
    "mixed_results2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ff_GB_results = mixed_results.append(mixed_results2)\n",
    "ff_GB_results.to_csv('/data/results/multimodal_pretrained_prediction_on_holdout_dataset.csv')"
   ]
  }
 ],
 "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}