{

 "cells": [

  {

   "cell_type": "code",

   "execution_count": 28,

   "metadata": {},

   "outputs": [],

   "source": [

    "#conda activate torch-xla-nightly\n",

    "#export XRT_TPU_CONFIG=\"tpu_worker;0;$10.0.101.2:8470\"\n",

    "#git init\n",

    "#git remote add origin https://github.com/nosound2/RSNA-Hemorrhage\n",

    "#git config remote.origin.push HEAD\n",

    "#git config credential.helper store\n",

    "#git pull origin master\n",

    "#git config --global branch.master.remote origin\n",

    "#git config --global branch.master.merge refs/heads/master\n",

    "#gcloud config set compute/zone europe-west4-a\n",

    "#gcloud config set compute/zone us-central1-c\n",

    "#gcloud auth login\n",

    "#gcloud config set project endless-empire-239015\n",

    "#pip install kaggle --user\n",

    "#PATH=$PATH:~/.local/bin\n",

    "#mkdir .kaggle\n",

    "#gsutil cp gs://recursion-double-strand/kaggle-keys/kaggle.json ~/.kaggle\n",

    "#chmod 600 /home/zahar_chikishev/.kaggle/kaggle.json\n",

    "#kaggle competitions download rsna-intracranial-hemorrhage-detection -f stage_2_train.csv\n",

    "#sudo apt install unzip\n",

    "#unzip stage_1_train.csv.zip\n",

    "#kaggle kernels output xhlulu/rsna-generate-metadata-csvs -p .\n",

    "#kaggle kernels output zaharch/dicom-test-metadata-to-csv -p .\n",

    "#kaggle kernels output zaharch/dicom-metadata-to-csv -p .\n",

    "#gsutil -m cp gs://rsna-hemorrhage/yuvals/* .\n",

    "\n",

    "#export XRT_TPU_CONFIG=\"tpu_worker;0;10.0.101.2:8470\"; conda activate torch-xla-nightly; jupyter notebook\n",

    "\n",

    "# 35.188.114.109\n",

    "\n",

    "# lsblk\n",

    "# sudo mkfs.ext4 -m 0 -F -E lazy_itable_init=0,lazy_journal_init=0,discard /dev/sdb\n",

    "# sudo mkdir /mnt/edisk\n",

    "# sudo mount -o discard,defaults /dev/sdb /mnt/edisk\n",

    "# sudo chmod a+w /mnt/edisk\n",

    "# df -H\n",

    "# echo UUID=`sudo blkid -s UUID -o value /dev/sdb` /mnt/edisk ext4 discard,defaults,nofail 0 2 | sudo tee -a /etc/fstab\n",

    "# cat /etc/fstab\n",

    "\n",

    "#jupyter notebook --generate-config\n",

    "#vi ~/.jupyter/jupyter_notebook_config.py\n",

    "\n",

    "#c = get_config()\n",

    "#c.NotebookApp.ip = '*'\n",

    "#c.NotebookApp.open_browser = False\n",

    "#c.NotebookApp.port = 5000"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": []

  },

  {

   "cell_type": "code",

   "execution_count": 2,

   "metadata": {},

   "outputs": [],

   "source": [

    "import sys\n",

    "\n",

    "from pathlib import Path\n",

    "from PIL import ImageDraw, ImageFont, Image\n",

    "from matplotlib import patches, patheffects\n",

    "import time\n",

    "from random import randint\n",

    "import numpy as np\n",

    "import pandas as pd\n",

    "import pickle\n",

    "\n",

    "from sklearn.model_selection import KFold, StratifiedKFold, GroupKFold\n",

    "from sklearn.preprocessing import LabelEncoder, OneHotEncoder, LabelBinarizer\n",

    "from sklearn.preprocessing import StandardScaler\n",

    "from sklearn.metrics import mean_squared_error,log_loss,roc_auc_score\n",

    "from scipy.stats import ks_2samp\n",

    "\n",

    "import pdb\n",

    "\n",

    "import scipy as sp\n",

    "from tqdm import tqdm, tqdm_notebook\n",

    "\n",

    "import os\n",

    "import glob\n",

    "\n",

    "import torch\n",

    "\n",

    "#CLOUD = not torch.cuda.is_available()\n",

    "CLOUD = (torch.cuda.get_device_name(0) in ['Tesla P4', 'Tesla K80', 'Tesla P100-PCIE-16GB'])\n",

    "CLOUD_SINGLE = True\n",

    "TPU = False\n",

    "\n",

    "if not CLOUD:\n",

    "    torch.cuda.current_device()\n",

    "\n",

    "import torch.nn as nn\n",

    "import torch.utils.data as D\n",

    "import torch.nn.functional as F\n",

    "import torch.utils as U\n",

    "\n",

    "import torchvision\n",

    "from torchvision import transforms as T\n",

    "from torchvision import models as M\n",

    "\n",

    "import matplotlib.pyplot as plt\n",

    "\n",

    "if CLOUD:\n",

    "    PATH = Path('/home/zahar_chikishev/Hemorrhage')\n",

    "    PATH_WORK = Path('/mnt/edisk/running')\n",

    "    PATH_DISK = Path('/mnt/edisk/running')\n",

    "else:\n",

    "    PATH = Path('C:/StudioProjects/Hemorrhage')\n",

    "    PATH_WORK = Path('C:/StudioProjects/Hemorrhage/running')\n",

    "    PATH_DISK = PATH_WORK\n",

    "\n",

    "from collections import defaultdict, Counter\n",

    "import random\n",

    "import seaborn as sn\n",

    "\n",

    "pd.set_option(\"display.max_columns\", 100)\n",

    "\n",

    "all_ich = ['any','epidural','intraparenchymal','intraventricular','subarachnoid','subdural']\n",

    "class_weights = 6.0*np.array([2,1,1,1,1,1])/7.0\n",

    "\n",

    "if CLOUD and TPU:\n",

    "    import torch_xla\n",

    "    import torch_xla.distributed.data_parallel as dp\n",

    "    import torch_xla.utils as xu\n",

    "    import torch_xla.core.xla_model as xm\n",

    "\n",

    "from typing import Collection"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 3,

   "metadata": {},

   "outputs": [],

   "source": [

    "all_black = '006d4432e'\n",

    "all_black = '00bd6c59c'\n",

    "\n",

    "if CLOUD:\n",

    "    if TPU:\n",

    "        device = xm.xla_device()\n",

    "    else:\n",

    "        device = 'cuda'\n",

    "    #device = 'cpu'\n",

    "    MAX_DEVICES = 1 if CLOUD_SINGLE else 8\n",

    "    NUM_WORKERS = 16\n",

    "    bs = 32\n",

    "else:\n",

    "    device = 'cuda'\n",

    "    #device = 'cpu'\n",

    "    MAX_DEVICES = 1\n",

    "    NUM_WORKERS = 0\n",

    "    bs = 16\n",

    "\n",

    "#if CLOUD and (not CLOUD_SINGLE):\n",

    "#    devices = xm.get_xla_supported_devices(max_devices=MAX_DEVICES)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 4,

   "metadata": {},

   "outputs": [],

   "source": [

    "SEED = 2351\n",

    "\n",

    "def setSeeds(seed):\n",

    "    np.random.seed(seed)\n",

    "    torch.manual_seed(seed)\n",

    "    torch.cuda.manual_seed(seed)\n",

    "\n",

    "setSeeds(SEED)\n",

    "torch.backends.cudnn.deterministic = True"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 5,

   "metadata": {},

   "outputs": [],

   "source": [

    "def getDSName(dataset):\n",

    "    if dataset == 6:\n",

    "        return 'Densenet201_F3'\n",

    "    elif dataset == 7:\n",

    "        return 'Densenet161_F3'\n",

    "    elif dataset == 8:\n",

    "        return 'Densenet169_F3'\n",

    "    elif dataset == 9:\n",

    "        return 'se_resnext101_32x4d_F3'\n",

    "    elif dataset == 10:\n",

    "        return 'se_resnet101_F3'\n",

    "    elif dataset == 11:\n",

    "        return 'se_resnext101_32x4d_F5'\n",

    "    elif dataset == 12:\n",

    "        return 'se_resnet101_F5'\n",

    "    elif dataset == 13:\n",

    "        return 'se_resnet101_focal_F5'\n",

    "    elif dataset == 14:\n",

    "        return 'se_resnet101_F5'\n",

    "    else: assert False"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 6,

   "metadata": {},

   "outputs": [],

   "source": [

    "def getDSParams(dataset):\n",

    "    if dataset == 6:\n",

    "        return 'Densenet201','_3','_v5',240,'','','classifier',''\n",

    "    elif dataset == 7:\n",

    "        return 'Densenet161','_3','_v4',552,'2','','classifier',''\n",

    "    elif dataset == 8:\n",

    "        return 'Densenet169','_3','_v3',208,'2','','classifier',''\n",

    "    elif dataset == 9:\n",

    "        return 'se_resnext101_32x4d','','',256,'','_tta','classifier',''\n",

    "    elif dataset == 10:\n",

    "        return 'se_resnet101','','',256,'','','classifier',''\n",

    "    elif dataset == 11:\n",

    "        return 'se_resnext101_32x4d_5','','',256,'','','new',''\n",

    "    elif dataset == 12:\n",

    "        return 'se_resnet101_5','','',256,'','','new',''\n",

    "    elif dataset == 13:\n",

    "        return 'se_resnet101_5f','','',256,'','','new','focal_'\n",

    "    elif dataset == 14:\n",

    "        return 'se_resnet101_5n','','',256,'','','new','stage2_'\n",

    "    else: assert False\n",

    "\n",

    "def getNFolds(dataset):\n",

    "    if dataset <= 10:\n",

    "        return 3\n",

    "    elif dataset <= 14:\n",

    "        return 5\n",

    "    else: assert False\n",

    "\n",

    "my_datasets3 = [7,9]\n",

    "my_datasets5 = [11,12,13]\n",

    "my_len = len(my_datasets3) + len(my_datasets5)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 7,

   "metadata": {},

   "outputs": [],

   "source": [

    "cols_le,cols_float,cols_bool = pickle.load(open(PATH_WORK/'covs','rb'))\n",

    "meta_cols = cols_bool + cols_float\n",

    "\n",

    "#meta_cols = [c for c in meta_cols if c not in ['SeriesPP']]\n",

    "#cols_le = cols_le[:-1]"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# Pre-processing"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 34,

   "metadata": {},

   "outputs": [],

   "source": [

    "def loadMetadata(rerun=False):\n",

    "    if rerun:\n",

    "        train_dedup = pd.read_csv(PATH_WORK/'yuval'/'train_dedup.csv')\n",

    "        pids, folding = pickle.load(open(PATH_WORK/'yuval'/'PID_splits.pkl','rb'))\n",

    "\n",

    "        assert len(pids) == 17079\n",

    "        assert len(np.unique(pids)) == 17079\n",

    "\n",

    "        for fol in folding:\n",

    "            assert len(fol[0]) + len(fol[1]) == 17079\n",

    "\n",

    "        assert len(folding[0][1]) + len(folding[1][1]) + len(folding[2][1]) == 17079\n",

    "\n",

    "        assert len(train_dedup.PID.unique()) == 17079\n",

    "\n",

    "        train_dedup['fold'] = np.nan\n",

    "\n",

    "        for fold in range(3):\n",

    "            train_dedup.loc[train_dedup.PID.isin(pids[folding[fold][1]]),'fold'] = fold\n",

    "\n",

    "        assert train_dedup.fold.isnull().sum() == 0\n",

    "        \n",

    "        oof5 = pickle.load(open(PATH_DISK/'yuval'/'OOF_validation_image_ids_5.pkl','rb'))\n",

    "        for fold in range(5):\n",

    "            train_dedup.loc[train_dedup.PatientID.isin(oof5[fold]),'fold5'] = fold\n",

    "        assert train_dedup.fold5.isnull().sum() == 0\n",

    "        \n",

    "        train_md = pd.read_csv(PATH_WORK/'train_md.csv').sort_values(['SeriesInstanceUID','pos_idx1'])\n",

    "        train_md['img_id'] = train_md.SOPInstanceUID.str.split('_').apply(lambda x: x[1])\n",

    "\n",

    "        ids_df = train_dedup[['fold','PatientID','fold5']]\n",

    "        ids_df.columns = ['fold','img_id','fold5']\n",

    "\n",

    "        train_md = ids_df.join(train_md.set_index('img_id'), on = 'img_id')\n",

    "\n",

    "        pickle.dump(train_md, open(PATH_WORK/'train.ids.df','wb'))\n",

    "\n",

    "        #test_md = pickle.load(open(PATH_WORK/'test.post.processed.1','rb'))\n",

    "        test_md = pd.read_csv(PATH_WORK/'test_md.csv').sort_values(['SeriesInstanceUID','pos_idx1'])\n",

    "        test_md['img_id'] = test_md.SOPInstanceUID.str.split('_').apply(lambda x: x[1])\n",

    "\n",

    "        filename = PATH_WORK/'yuval'/'test_indexes.pkl'\n",

    "        test_ids = pickle.load(open(filename,'rb'))\n",

    "\n",

    "        test_ids_df = pd.DataFrame(test_ids, columns = ['img_id'])\n",

    "        test_md = test_ids_df.join(test_md.set_index('img_id'), on = 'img_id')\n",

    "        \n",

    "        assert len(test_md.SeriesInstanceUID.unique()) == 2214\n",

    "        \n",

    "        #train_md = pd.concat([train_md, test_md], axis = 0, sort=False).reset_index(drop=True)\n",

    "        \n",

    "        #pickle.dump(train_md, open(PATH_WORK/'train.ids.df','wb'))\n",

    "        pickle.dump(test_md, open(PATH_WORK/'test.ids.df','wb'))\n",

    "    else:\n",

    "        train_md = pickle.load(open(PATH_WORK/'train.ids.df','rb'))\n",

    "        test_md = pickle.load(open(PATH_WORK/'test.ids.df','rb'))\n",

    "    \n",

    "    return train_md, test_md"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 17,

   "metadata": {},

   "outputs": [],

   "source": [

    "def loadMetadata2(rerun=False):\n",

    "    if rerun:\n",

    "        train_dedup2 = pd.read_csv(PATH_WORK/'yuval'/'train_stage2.csv')\n",

    "\n",

    "        assert train_dedup2.PatientID.nunique() == 752797\n",

    "        assert train_dedup2.PID.nunique() == 18938\n",

    "\n",

    "        pids, folding = pickle.load(open(PATH_WORK/'yuval'/'PID_splits.pkl','rb'))\n",

    "        pids3, folding3 = pickle.load(open(PATH_WORK/'yuval'/'PID_splits_3_stage_2.pkl','rb'))\n",

    "        pids5, folding5 = pickle.load(open(PATH_WORK/'yuval'/'PID_splits_5_stage_2.pkl','rb'))\n",

    "\n",

    "        assert np.array([len(folding3[i][1]) for i in range(3)]).sum() == 18938\n",

    "        assert np.array([len(folding3[i][0]) for i in range(3)]).sum() == 2*18938\n",

    "        \n",

    "        for i in range(3):\n",

    "            assert set(pids[folding[i][1]]).issubset(set(pids3[folding3[i][1]]))\n",

    "\n",

    "        for i in range(3):\n",

    "            train_dedup2.loc[train_dedup2.PID.isin(pids3[folding3[i][1]]),'fold'] = i\n",

    "        assert train_dedup2.fold.isnull().sum() == 0\n",

    "\n",

    "        for i in range(5):\n",

    "            train_dedup2.loc[train_dedup2.PID.isin(pids5[folding5[i][1]]),'fold5'] = i\n",

    "        assert train_dedup2.fold5.isnull().sum() == 0\n",

    "        \n",

    "        oof5 = pickle.load(open(PATH_DISK/'yuval'/'OOF_validation_image_ids_5.pkl','rb'))\n",

    "        for i in range(5):\n",

    "            assert set(oof5[i]).issubset(set(train_dedup2.loc[train_dedup2.fold5==i,'PatientID']))\n",

    "        \n",

    "        train_csv = pd.read_csv(PATH/'stage_2_train.csv')\n",

    "        train_csv = train_csv.loc[~train_csv.ID.duplicated()].sort_values('ID').reset_index(drop=True)        \n",

    "        all_sop_ids = train_csv.ID.str.split('_').apply(lambda x: x[0]+'_'+x[1]).unique()\n",

    "        train_df = pd.DataFrame(train_csv.Label.values.reshape((-1,6)), columns = all_ich)\n",

    "        train_df['sop_id'] = all_sop_ids\n",

    "        \n",

    "        old_test = pd.read_csv(PATH_WORK/'test_md.csv')\n",

    "        old_test = old_test.drop(all_ich,axis=1)\n",

    "        old_test = old_test.join(train_df.set_index('sop_id'), on = 'SOPInstanceUID')\n",

    "        \n",

    "        train_md2 = pd.concat([pd.read_csv(PATH_WORK/'train_md.csv'), old_test],\\\n",

    "                              axis=0,sort=False)\n",

    "        train_md2['img_id'] = train_md2.SOPInstanceUID.str.split('_').apply(lambda x: x[1])\n",

    "        \n",

    "        ids_df = train_dedup2[['fold','PatientID','fold5']]\n",

    "        ids_df.columns = ['fold','img_id','fold5']\n",

    "\n",

    "        train_md2 = ids_df.join(train_md2.set_index('img_id'), on = 'img_id')\n",

    "        assert train_md2.test.sum() == 78545\n",

    "        \n",

    "        pickle.dump(train_md2, open(PATH_WORK/'train2.ids.df','wb'))\n",

    "        \n",

    "        test_md = pd.read_csv(PATH_WORK/'test2_md.csv').sort_values(['SeriesInstanceUID','pos_idx1'])\n",

    "        test_md['img_id'] = test_md.SOPInstanceUID.str.split('_').apply(lambda x: x[1])\n",

    "\n",

    "        filename = PATH_WORK/'yuval/test_indexes_stage2.pkl'\n",

    "        test_ids = pickle.load(open(filename,'rb'))\n",

    "\n",

    "        test_ids_df = pd.DataFrame(test_ids, columns = ['img_id'])\n",

    "        test_md = test_ids_df.join(test_md.set_index('img_id'), on = 'img_id')\n",

    "\n",

    "        assert len(test_md.SeriesInstanceUID.unique()) == 3518\n",

    "\n",

    "        pickle.dump(test_md, open(PATH_WORK/'test2.ids.df','wb'))\n",

    "\n",

    "    else:\n",

    "        train_md2 = pickle.load(open(PATH_WORK/'train2.ids.df','rb'))\n",

    "        test_md = pickle.load(open(PATH_WORK/'test2.ids.df','rb'))\n",

    "    \n",

    "    return train_md2, test_md"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 34,

   "metadata": {},

   "outputs": [],

   "source": [

    "def loadMetadata3(rerun=False):\n",

    "    if rerun:\n",

    "        train_md = pickle.load(open(PATH_WORK/'train.ids.df','rb'))\n",

    "        test_md = pickle.load(open(PATH_WORK/'test.ids.df','rb'))\n",

    "        train_md = pd.concat([train_md, test_md], axis=0,sort=False).reset_index(drop=True)\n",

    "        \n",

    "        train_md.fold = np.nan\n",

    "        train_md['PID'] = train_md.PatientID.str.split('_').apply(lambda x: x[1])\n",

    "        \n",

    "        pids, folding = pickle.load(open(PATH_WORK/'yuval'/'PID_splits.pkl','rb'))\n",

    "        pids3, folding3 = pickle.load(open(PATH_WORK/'yuval'/'PID_splits_3_stage_2.pkl','rb'))\n",

    "        pids5, folding5 = pickle.load(open(PATH_WORK/'yuval'/'PID_splits_5_stage_2.pkl','rb'))\n",

    "\n",

    "        for i in range(3):\n",

    "            train_md.loc[train_md.PID.isin(pids3[folding3[i][1]]),'fold'] = i\n",

    "        assert train_md.fold.isnull().sum() == 0\n",

    "\n",

    "        for i in range(5):\n",

    "            train_md.loc[train_md.PID.isin(pids5[folding5[i][1]]),'fold5'] = i\n",

    "        assert train_md.fold5.isnull().sum() == 0\n",

    "        \n",

    "        train_md.weights = train_md.weights.fillna(1)\n",

    "        \n",

    "        pickle.dump(train_md, open(PATH_WORK/'train3.ids.df','wb'))\n",

    "        \n",

    "        test_md = pd.read_csv(PATH_WORK/'test2_md.csv').sort_values(['SeriesInstanceUID','pos_idx1'])\n",

    "        test_md['img_id'] = test_md.SOPInstanceUID.str.split('_').apply(lambda x: x[1])\n",

    "\n",

    "        filename = PATH_WORK/'yuval/test_indexes_stage2.pkl'\n",

    "        test_ids = pickle.load(open(filename,'rb'))\n",

    "\n",

    "        test_ids_df = pd.DataFrame(test_ids, columns = ['img_id'])\n",

    "        test_md = test_ids_df.join(test_md.set_index('img_id'), on = 'img_id')\n",

    "\n",

    "        assert len(test_md.SeriesInstanceUID.unique()) == 3518\n",

    "\n",

    "        pickle.dump(test_md, open(PATH_WORK/'test3.ids.df','wb'))\n",

    "    else:\n",

    "        train_md = pickle.load(open(PATH_WORK/'train3.ids.df','rb'))\n",

    "        test_md = pickle.load(open(PATH_WORK/'test3.ids.df','rb'))\n",

    "    \n",

    "    return train_md, test_md"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 66,

   "metadata": {},

   "outputs": [],

   "source": [

    "#train_md, test_md = loadMetadata(True)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 18,

   "metadata": {},

   "outputs": [],

   "source": [

    "#train_md = loadMetadata2(True)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": []

  },

  {

   "cell_type": "code",

   "execution_count": 7,

   "metadata": {},

   "outputs": [],

   "source": [

    "def preprocessedData(dataset, folds = range(3), fold_col=None, do_test=True, do_test2=True, do_train=True):\n",

    "    assert dataset >= 6\n",

    "    \n",

    "    dataset_name, filename_add, filename_add2, feat_sz, ds_num, test_fix, dsft, focal = getDSParams(dataset)\n",

    "    \n",

    "    #if 'train_md' not in globals() or 'test_md' not in globals():\n",

    "    #    train_md, test_md = loadMetadata(False)\n",

    "    \n",

    "    PATH_DS = PATH_DISK/'features/{}{}'.format(dataset_name,filename_add2)\n",

    "    if not PATH_DS.is_dir():\n",

    "        PATH_DS.mkdir()\n",

    "    if not (PATH_DS/'train').is_dir():\n",

    "        (PATH_DS/'train').mkdir()\n",

    "    if not (PATH_DS/'test').is_dir():\n",

    "        (PATH_DS/'test').mkdir()\n",

    "    if not (PATH_DS/'test2').is_dir():\n",

    "        (PATH_DS/'test2').mkdir()\n",

    "    \n",

    "    if fold_col is None:\n",

    "        fold_col = 'fold'\n",

    "        if len(folds) == 5:\n",

    "            fold_col = 'fold5'\n",

    "    \n",

    "    for fold in folds:\n",

    "        \n",

    "        if do_train:\n",

    "            filename = PATH_DISK/'yuval'/\\\n",

    "                'model_{}{}_version_{}_splits_{}type_features_train_tta{}_split_{}.pkl'\\\n",

    "                .format(dataset_name.replace('_5n','').replace('_5f','').replace('_5',''), \\\n",

    "                        filename_add, dsft, focal, ds_num, fold)\n",

    "            feats = pickle.load(open(filename,'rb'))\n",

    "\n",

    "            print('dataset',dataset,'fold',fold,'feats size', feats.shape)\n",

    "            if dataset <= 13:\n",

    "                train_md_loc = train_md.loc[~train_md.test]\n",

    "            else:\n",

    "                train_md_loc = train_md\n",

    "            \n",

    "            assert len(feats) == 4*len(train_md_loc)\n",

    "            assert feats.shape[1] == feat_sz\n",

    "            means = feats.mean(0,keepdim=True)\n",

    "            stds = feats.std(0,keepdim=True)\n",

    "\n",

    "            feats = feats - means\n",

    "            feats = torch.where(stds > 0, feats/stds, feats)\n",

    "\n",

    "            for i in range(4):\n",

    "                feats_sub1 = feats[torch.BoolTensor(np.arange(len(feats))%4 == i)]\n",

    "                feats_sub2 = feats_sub1[torch.BoolTensor(train_md_loc[fold_col] != fold)]\n",

    "                pickle.dump(feats_sub2, open(PATH_DS/'train/train.f{}.a{}'.format(fold,i),'wb'))\n",

    "\n",

    "                feats_sub2 = feats_sub1[torch.BoolTensor(train_md_loc[fold_col] == fold)]\n",

    "                pickle.dump(feats_sub2, open(PATH_DS/'train/valid.f{}.a{}'.format(fold,i),'wb'))\n",

    "\n",

    "                if i==0:\n",

    "                    black_feats = feats_sub1[torch.BoolTensor(train_md_loc.img_id == all_black)].squeeze()\n",

    "                    pickle.dump(black_feats, open(PATH_DS/'train/black.f{}'.format(fold),'wb'))\n",

    "        \n",

    "        if do_test:\n",

    "            filename = PATH_DISK/'yuval'/\\\n",

    "                'model_{}{}_version_{}_splits_{}type_features_test{}{}_split_{}.pkl'\\\n",

    "                .format(dataset_name.replace('_5n','').replace('_5f','').replace('_5',''),\n",

    "                        filename_add,dsft,focal,ds_num,test_fix,fold)\n",

    "            feats = pickle.load(open(filename,'rb'))\n",

    "\n",

    "            assert len(feats) == 8*len(test_md)\n",

    "            assert feats.shape[1] == feat_sz\n",

    "\n",

    "            feats = feats - means\n",

    "            feats = torch.where(stds > 0, feats/stds, feats)\n",

    "\n",

    "            for i in range(8):\n",

    "                feats_sub = feats[torch.BoolTensor(np.arange(len(feats))%8 == i)]\n",

    "                pickle.dump(feats_sub, open(PATH_DS/'test/test.f{}.a{}'.format(fold,i),'wb'))\n",

    "                assert len(feats_sub) == len(test_md)\n",

    "        \n",

    "        if do_test2:\n",

    "            filename = PATH_DISK/'yuval'/\\\n",

    "                'model_{}{}_version_{}_splits_{}type_features_test_stage2_split_{}.pkl'\\\n",

    "                .format(dataset_name.replace('_5n','').replace('_5f','').replace('_5',''),\n",

    "                        filename_add,dsft,focal,fold)\n",

    "            feats = pickle.load(open(filename,'rb'))\n",

    "\n",

    "            assert len(feats) == 8*len(test_md)\n",

    "            assert feats.shape[1] == feat_sz\n",

    "\n",

    "            feats = feats - means\n",

    "            feats = torch.where(stds > 0, feats/stds, feats)\n",

    "\n",

    "            for i in range(8):\n",

    "                feats_sub = feats[torch.BoolTensor(np.arange(len(feats))%8 == i)]\n",

    "                pickle.dump(feats_sub, open(PATH_DS/'test2/test.f{}.a{}'.format(fold,i),'wb'))\n",

    "                assert len(feats_sub) == len(test_md)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": []

  },

  {

   "cell_type": "code",

   "execution_count": null,

   "metadata": {},

   "outputs": [],

   "source": []

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# Dataset"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 86,

   "metadata": {},

   "outputs": [],

   "source": [

    "if False:\n",

    "    path = PATH_WORK/'features/densenet161_v3/train/ID_992b567eb6'\n",

    "    black_feats = pickle.load(open(path,'rb'))[41]"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 87,

   "metadata": {},

   "outputs": [],

   "source": [

    "class RSNA_DataSet(D.Dataset):\n",

    "    def __init__(self, metadata, dataset, mode='train', bs=None, fold=0):\n",

    "        \n",

    "        super(RSNA_DataSet, self).__init__()\n",

    "        \n",

    "        dataset_name, filename_add, filename_add2, feat_sz,_,_,_,_ = getDSParams(dataset)\n",

    "        num_folds = getNFolds(dataset)\n",

    "        \n",

    "        self.dataset_name = dataset_name\n",

    "        self.filename_add2 = filename_add2\n",

    "        \n",

    "        folds_col = 'fold'\n",

    "        if num_folds == 5:\n",

    "            folds_col = 'fold5'\n",

    "        \n",

    "        if (dataset <= 13) and (mode in ['train','valid']) and TRAIN_ON_STAGE_1:\n",

    "            if mode == 'train':\n",

    "                self.test_mask = torch.BoolTensor((metadata.loc[metadata.test, folds_col] != fold).values)\n",

    "            else:\n",

    "                self.test_mask = torch.BoolTensor((metadata.loc[metadata.test, folds_col] == fold).values)\n",

    "        \n",

    "        if mode == 'train':\n",

    "            md = metadata.loc[metadata[folds_col] != fold].copy().reset_index(drop=True)\n",

    "        elif mode == 'valid':\n",

    "            md = metadata.loc[metadata[folds_col] == fold].copy().reset_index(drop=True)\n",

    "        else:\n",

    "            md = metadata.copy().reset_index(drop=True)\n",

    "        \n",

    "        series = np.sort(md.SeriesInstanceUID.unique())\n",

    "        md = md.set_index('SeriesInstanceUID', drop=True)\n",

    "        \n",

    "        samples_add = 0\n",

    "        if (mode != 'train') and not DATA_SMALL:\n",

    "            batch_num = -((-len(series))//(bs*MAX_DEVICES))\n",

    "            samples_add = batch_num*bs*MAX_DEVICES - len(series)\n",

    "            print('adding dummy serieses', samples_add)\n",

    "        \n",

    "        #self.records = df.to_records(index=False)\n",

    "        self.mode = mode\n",

    "        self.real = np.concatenate([np.repeat(True,len(series)),np.repeat(False,samples_add)])\n",

    "        self.series = np.concatenate([series, random.sample(list(series),samples_add)])\n",

    "        self.metadata = md\n",

    "        self.dataset = dataset\n",

    "        self.fold = fold\n",

    "        \n",

    "        print('DataSet', dataset, mode, 'size', len(self.series), 'fold', fold)\n",

    "        \n",

    "        path = PATH_DISK/'features/{}{}/train/black.f{}'.format(self.dataset_name, self.filename_add2, fold)\n",

    "        self.black_feats = pickle.load(open(path,'rb')).squeeze()\n",

    "        \n",

    "        if WEIGHTED and (self.mode == 'train'):\n",

    "            tt = self.metadata['weights'].groupby(self.metadata.index).mean()\n",

    "            self.weights = tt.loc[self.series].values\n",

    "            print(pd.value_counts(self.weights))\n",

    "    \n",

    "    def setFeats(self, anum, epoch=0):\n",

    "        def getAPath(an,mode):\n",

    "            folder = 'test2' if mode == 'test' else 'train'\n",

    "            return PATH_DISK/'features/{}{}/{}/{}.f{}.a{}'\\\n",

    "                .format(self.dataset_name,self.filename_add2,folder,mode,self.fold,an)\n",

    "        \n",

    "        def getAPathFeats(mode, sz, mask=None):\n",

    "            max_a = 8 if mode == 'test' else 4\n",

    "            feats2 = torch.stack([pickle.load(open(getAPath(an,mode),'rb')) for an in range(max_a)])\n",

    "            if mask is not None:\n",

    "                feats2 = feats2[:,mask]\n",

    "            assert feats2.shape[1] == sz\n",

    "            feats = feats2[torch.randint(max_a,(sz,)), torch.arange(sz, dtype=torch.long)].squeeze()\n",

    "            return feats\n",

    "        \n",

    "        if self.dataset == 1: return\n",

    "        print('setFeats, augmentation', anum)\n",

    "        self.anum = anum\n",

    "        sz = len(self.metadata)\n",

    "        \n",

    "        assert anum <= 0\n",

    "        \n",

    "        if anum == -1:\n",

    "            if (self.dataset <= 13) and (self.mode in ['train','valid']) and TRAIN_ON_STAGE_1:\n",

    "                feats = torch.cat([getAPathFeats(self.mode, sz - self.metadata.test.sum()), \n",

    "                                   getAPathFeats('test', self.metadata.test.sum(), self.test_mask)] ,axis=0)\n",

    "            else:\n",

    "                feats = getAPathFeats(self.mode, sz)\n",

    "        else:\n",

    "            if (self.dataset <= 13) and (self.mode in ['train','valid']) and TRAIN_ON_STAGE_1:\n",

    "                feats = torch.cat([pickle.load(open(getAPath(anum,self.mode),'rb')),\n",

    "                                   pickle.load(open(getAPath(anum,'test'),'rb'))[self.test_mask]], axis=0)\n",

    "            else:\n",

    "                feats = pickle.load(open(getAPath(anum,self.mode),'rb'))\n",

    "        \n",

    "        self.feats = feats\n",

    "        self.epoch = epoch\n",

    "        assert len(feats) == sz\n",

    "    \n",

    "    def __getitem__(self, index):\n",

    "        \n",

    "        series_id = self.series[index]\n",

    "        #df = self.metadata.loc[self.metadata.SeriesInstanceUID == series_id].reset_index(drop=True)\n",

    "        df = self.metadata.loc[series_id].reset_index(drop=True)\n",

    "        \n",

    "        if self.dataset >= 6:\n",

    "            feats = self.feats[torch.BoolTensor(self.metadata.index.values == series_id)]\n",

    "        else: assert False\n",

    "        \n",

    "        order = np.argsort(df.pos_idx1.values)\n",

    "        df = df.sort_values(['pos_idx1'])\n",

    "        feats = feats[torch.LongTensor(order)]\n",

    "        \n",

    "        #if WEIGHTED:\n",

    "        #    non_black = torch.Tensor(df['weights'])\n",

    "        #else:\n",

    "        non_black = torch.ones(len(feats))\n",

    "        \n",

    "        feats = torch.cat([feats, torch.Tensor(df[meta_cols].values)], dim=1)\n",

    "        feats_le = torch.LongTensor(df[cols_le].values)\n",

    "        \n",

    "        target = torch.Tensor(df[all_ich].values)\n",

    "        \n",

    "        PAD = 12\n",

    "        \n",

    "        np.random.seed(1234 + index + 10000*self.epoch)\n",

    "        offset = np.random.randint(0, 61 - feats.shape[0])\n",

    "        \n",

    "        #offset = 0\n",

    "        bk_add = 0\n",

    "        top_pad = PAD + offset\n",

    "        if top_pad > 0:\n",

    "            dummy_row = torch.cat([self.black_feats, torch.Tensor(df.head(1)[meta_cols].values).squeeze()])\n",

    "            feats = torch.cat([dummy_row.repeat(top_pad,1), feats], dim=0)\n",

    "            feats_le = torch.cat([torch.LongTensor(df.head(1)[cols_le].values).squeeze().repeat(top_pad,1), feats_le])\n",

    "            if offset > 0:\n",

    "                non_black = torch.cat([bk_add + torch.zeros(offset), non_black])\n",

    "                target = torch.cat([torch.zeros((offset, len(all_ich))), target], dim=0)\n",

    "        bot_pad = 60 - len(df) - offset + PAD\n",

    "        if bot_pad > 0:\n",

    "            dummy_row = torch.cat([self.black_feats, torch.Tensor(df.tail(1)[meta_cols].values).squeeze()])\n",

    "            feats = torch.cat([feats, dummy_row.repeat(bot_pad,1)], dim=0)\n",

    "            feats_le = torch.cat([feats_le, torch.LongTensor(df.tail(1)[cols_le].values).squeeze().repeat(bot_pad,1)])\n",

    "            if (60 - len(df) - offset) > 0:\n",

    "                non_black = torch.cat([non_black, bk_add + torch.zeros(60 - len(df) - offset)])\n",

    "                target = torch.cat([target, torch.zeros((60 - len(df) - offset, len(all_ich)))], dim=0)\n",

    "        \n",

    "        assert feats_le.shape[0] == (60 + 2*PAD)\n",

    "        assert feats.shape[0] == (60 + 2*PAD)\n",

    "        assert target.shape[0] == 60\n",

    "        \n",

    "        idx = index\n",

    "        if not self.real[index]: idx = -1\n",

    "        \n",

    "        if self.mode == 'train':\n",

    "            return feats, feats_le, target, non_black\n",

    "        else:\n",

    "            return feats, feats_le, target, idx, offset\n",

    "    \n",

    "    def __len__(self):\n",

    "        return len(self.series) if not DATA_SMALL else int(0.01*len(self.series))"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 88,

   "metadata": {},

   "outputs": [],

   "source": [

    "def getCurrentBatch(dataset, fold=0, ver=VERSION):\n",

    "    sel_batch = None\n",

    "    for filename in os.listdir(PATH_DISK/'models'):\n",

    "        splits = filename.split('.')\n",

    "        if int(splits[2][1]) != fold: continue\n",

    "        if int(splits[3][1:]) != dataset: continue\n",

    "        if int(splits[4][1:]) != ver: continue\n",

    "        if sel_batch is None:\n",

    "            sel_batch = int(splits[1][1:])\n",

    "        else:\n",

    "            sel_batch = max(sel_batch, int(splits[1][1:]))\n",

    "    return sel_batch\n",

    "\n",

    "def modelFileName(dataset, fold=0, batch = 1, return_last = False, return_next = False, ver=VERSION):\n",

    "    sel_batch = batch\n",

    "    if return_last or return_next:\n",

    "        sel_batch = getCurrentBatch(fold=fold, dataset=dataset, ver=ver)\n",

    "        if return_last and sel_batch is None:\n",

    "            return None\n",

    "        if return_next:\n",

    "            if sel_batch is None: sel_batch = 1\n",

    "            else: sel_batch += 1\n",

    "    \n",

    "    return 'model.b{}.f{}.d{}.v{}'.format(sel_batch, fold, dataset, ver)"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# Model"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 6,

   "metadata": {},

   "outputs": [],

   "source": [

    "class BCEWithLogitsLoss(nn.Module):\n",

    "    def __init__(self, weight=None):\n",

    "        super().__init__()\n",

    "        self.weight = weight\n",

    "    \n",

    "    def forward(self, input, target, batch_weights = None, focal=0):\n",

    "        if focal == 0:\n",

    "            loss = (torch.log(1+torch.exp(input)) - target*input)*self.weight\n",

    "        else:\n",

    "            loss = torch.pow(1+torch.exp(input), -focal) * \\\n",

    "                (((1-target)*torch.exp(focal*input) + target) * torch.log(1+torch.exp(input)) - target*input)\n",

    "            loss = loss*self.weight\n",

    "        if batch_weights is not None:\n",

    "            loss = batch_weights*loss\n",

    "        return loss.mean()"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 90,

   "metadata": {},

   "outputs": [],

   "source": [

    "def bn_drop_lin(n_in:int, n_out:int, bn:bool=True, p:float=0., actn=None):\n",

    "    \"Sequence of batchnorm (if `bn`), dropout (with `p`) and linear (`n_in`,`n_out`) layers followed by `actn`.\"\n",

    "    layers = [nn.BatchNorm1d(n_in)] if bn else []\n",

    "    if p != 0: layers.append(nn.Dropout(p))\n",

    "    layers.append(nn.Linear(n_in, n_out))\n",

    "    if actn is not None: layers.append(actn)\n",

    "    return layers"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 91,

   "metadata": {},

   "outputs": [],

   "source": [

    "def noop(x): return x\n",

    "act_fun = nn.ReLU(inplace=True)\n",

    "\n",

    "def conv_layer(ni, nf, ks=3, act=True):\n",

    "    bn = nn.BatchNorm1d(nf)\n",

    "    layers = [nn.Conv1d(ni, nf, ks), bn]\n",

    "    if act: layers.append(act_fun)\n",

    "    return nn.Sequential(*layers)\n",

    "\n",

    "class ResBlock(nn.Module):\n",

    "    def __init__(self, ni, nh):\n",

    "        super().__init__()\n",

    "        layers  = [conv_layer(ni, nh, 1),\n",

    "                   conv_layer(nh, nh, 5, act=False)]\n",

    "        self.convs = nn.Sequential(*layers)\n",

    "        self.idconv = noop if (ni == nh) else conv_layer(ni, nh, 1, act=False)\n",

    "    \n",

    "    def forward(self, x): return act_fun(self.convs(x) + self.idconv(x[:,:,2:-2]))"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 77,

   "metadata": {},

   "outputs": [],

   "source": [

    "class ResNetModel(nn.Module):\n",

    "    def __init__(self, n_cont:int, feat_sz=2208):\n",

    "        super().__init__()\n",

    "        \n",

    "        self.le_sz = 10\n",

    "        assert self.le_sz == (len(cols_le))\n",

    "        le_in_sizes = np.array([5,5,7,4,4,11,4,6,3,3])\n",

    "        le_out_sizes = np.array([3,3,4,2,2,6,2,4,2,2])\n",

    "        le_out_sz = le_out_sizes.sum()\n",

    "        self.embeddings = nn.ModuleList([embedding(le_in_sizes[i], le_out_sizes[i]) for i in range(self.le_sz)])\n",

    "        \n",

    "        self.feat_sz = feat_sz\n",

    "        \n",

    "        self.n_cont = n_cont\n",

    "        scale = 4\n",

    "        \n",

    "        self.conv2D = nn.Conv2d(1,scale*16,(feat_sz + n_cont + le_out_sz,1))\n",

    "        self.bn1 = nn.BatchNorm1d(scale*16)\n",

    "        \n",

    "        self.res1 = ResBlock(scale*16,scale*16)\n",

    "        self.res2 = ResBlock(scale*16,scale*8)\n",

    "        \n",

    "        self.res3 = ResBlock(scale*24,scale*16)\n",

    "        self.res4 = ResBlock(scale*16,scale*8)\n",

    "        \n",

    "        self.res5 = ResBlock(scale*32,scale*16)\n",

    "        self.res6 = ResBlock(scale*16,scale*8)\n",

    "        \n",

    "        self.conv1D = nn.Conv1d(scale*40,6,1)\n",

    "    \n",

    "    def forward(self, x, x_le, x_le_mix = None, lambd = None) -> torch.Tensor:\n",

    "        x_le = [e(x_le[:,:,i]) for i,e in enumerate(self.embeddings)]\n",

    "        x_le = torch.cat(x_le, 2)\n",

    "        \n",

    "        x = torch.cat([x, x_le], 2)\n",

    "        x = x.transpose(1,2)\n",

    "        \n",

    "        #x = torch.cat([x[:,:self.feat_sz],self.bn_cont(x[:,self.feat_sz:])], dim=1)\n",

    "        x = x.reshape(x.shape[0],1,x.shape[1],x.shape[2])\n",

    "        \n",

    "        x = self.conv2D(x).squeeze()\n",

    "        x = self.bn1(x)\n",

    "        x = act_fun(x)\n",

    "        \n",

    "        x2 = self.res1(x)\n",

    "        x2 = self.res2(x2)\n",

    "        \n",

    "        x3 = torch.cat([x[:,:,4:-4], x2], 1)\n",

    "        \n",

    "        x3 = self.res3(x3)\n",

    "        x3 = self.res4(x3)\n",

    "        \n",

    "        x4 = torch.cat([x[:,:,8:-8], x2[:,:,4:-4], x3], 1)\n",

    "        \n",

    "        x4 = self.res5(x4)\n",

    "        x4 = self.res6(x4)\n",

    "        \n",

    "        x5 = torch.cat([x[:,:,12:-12], x2[:,:,8:-8], x3[:,:,4:-4], x4], 1)\n",

    "        x5 = self.conv1D(x5)\n",

    "        x5 = x5.transpose(1,2)\n",

    "        \n",

    "        return x5"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 93,

   "metadata": {},

   "outputs": [],

   "source": [

    "def trunc_normal_(x, mean:float=0., std:float=1.):\n",

    "    \"Truncated normal initialization.\"\n",

    "    # From https://discuss.pytorch.org/t/implementing-truncated-normal-initializer/4778/12\n",

    "    return x.normal_().fmod_(2).mul_(std).add_(mean)"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 94,

   "metadata": {},

   "outputs": [],

   "source": [

    "def embedding(ni:int,nf:int) -> nn.Module:\n",

    "    \"Create an embedding layer.\"\n",

    "    emb = nn.Embedding(ni, nf)\n",

    "    # See https://arxiv.org/abs/1711.09160\n",

    "    with torch.no_grad(): trunc_normal_(emb.weight, std=0.01)\n",

    "    return emb"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 95,

   "metadata": {},

   "outputs": [],

   "source": [

    "class TabularModel(nn.Module):\n",

    "    \"Basic model for tabular data.\"\n",

    "    def __init__(self, n_cont:int, feat_sz=2208, fc_drop_p=0.3):\n",

    "        super().__init__()\n",

    "        self.le_sz = 10\n",

    "        assert len(cols_le) == self.le_sz\n",

    "        le_in_sizes = np.array([5,5,7,4,4,11,4,6,3,3])\n",

    "        le_out_sizes = np.array([3,3,4,2,2,6,2,4,2,2])\n",

    "        le_out_sz = le_out_sizes.sum()\n",

    "        self.embeddings = nn.ModuleList([embedding(le_in_sizes[i], le_out_sizes[i]) for i in range(self.le_sz)])\n",

    "        #self.bn_cont = nn.BatchNorm1d(feat_sz + n_cont)\n",

    "        \n",

    "        self.feat_sz = feat_sz\n",

    "        self.bn_cont = nn.BatchNorm1d(n_cont + le_out_sz)\n",

    "        self.n_cont = n_cont\n",

    "        self.fc_drop = nn.Dropout(p=fc_drop_p)\n",

    "        self.relu = nn.ReLU(inplace=True)\n",

    "        \n",

    "        scale = 4\n",

    "        \n",

    "        self.conv2D_1 = nn.Conv2d(1,16*scale,(feat_sz + n_cont + le_out_sz,1))\n",

    "        self.conv2D_2 = nn.Conv2d(1,16*scale,(feat_sz + n_cont + le_out_sz,5))\n",

    "        self.bn_cont1 = nn.BatchNorm1d(32*scale)\n",

    "        self.conv1D_1 = nn.Conv1d(32*scale,16*scale,3)\n",

    "        self.conv1D_3 = nn.Conv1d(32*scale,16*scale,5,dilation=5)\n",

    "        self.conv1D_2 = nn.Conv1d(32*scale,6,3)\n",

    "        self.bn_cont2 = nn.BatchNorm1d(32*scale)\n",

    "        self.bn_cont3 = nn.BatchNorm1d(6)\n",

    "\n",

    "        self.conv1D_4 = nn.Conv1d(32*scale,32*scale,3)\n",

    "        self.bn_cont4 = nn.BatchNorm1d(32*scale)\n",

    "\n",

    "    def forward(self, x, x_le, x_le_mix = None, lambd = None) -> torch.Tensor:\n",

    "        x_le = [e(x_le[:,:,i]) for i,e in enumerate(self.embeddings)]\n",

    "        x_le = torch.cat(x_le, 2)\n",

    "        \n",

    "        if MIXUP and x_le_mix is not None:\n",

    "            x_le_mix = [e(x_le_mix[:,:,i]) for i,e in enumerate(self.embeddings)]\n",

    "            x_le_mix = torch.cat(x_le_mix, 2)\n",

    "            x_le = lambd * x_le + (1-lambd) * x_le_mix\n",

    "        \n",

    "        #assert torch.isnan(x_le).any().cpu() == False\n",

    "        x = torch.cat([x, x_le], 2)\n",

    "        x = x.transpose(1,2)\n",

    "        \n",

    "        #x = torch.cat([x[:,:self.feat_sz],self.bn_cont(x[:,self.feat_sz:])], dim=1)\n",

    "        \n",

    "        x = x.reshape(x.shape[0],1,x.shape[1],x.shape[2])\n",

    "        x = self.fc_drop(x)\n",

    "        x = torch.cat([self.conv2D_1(x[:,:,:,2:(-2)]).squeeze(), \n",

    "                       self.conv2D_2(x).squeeze()], dim=1)\n",

    "        x = self.relu(x)\n",

    "        x = self.bn_cont1(x)\n",

    "        x = self.fc_drop(x)\n",

    "        \n",

    "        x = torch.cat([self.conv1D_1(x[:,:,9:(-9)]),\n",

    "                       self.conv1D_3(x)], dim=1)\n",

    "        x = self.relu(x)\n",

    "        x = self.bn_cont2(x)\n",

    "        x = self.fc_drop(x)\n",

    "        \n",

    "        x = self.conv1D_4(x)\n",

    "        x = self.relu(x)\n",

    "        x = self.bn_cont4(x)\n",

    "        \n",

    "        x = self.conv1D_2(x)\n",

    "        x = x.transpose(1,2)\n",

    "        \n",

    "        return x"

   ]

  },

  {

   "cell_type": "markdown",

   "metadata": {},

   "source": [

    "# Training"

   ]

  },

  {

   "cell_type": "code",

   "execution_count": 96,

   "metadata": {},

   "outputs": [],

   "source": [

    "def train_loop_fn(model, loader, device, context = None):\n",

    "    \n",

    "    if CLOUD and (not CLOUD_SINGLE):\n",

    "        tlen = len(loader._loader._loader)\n",

    "        OUT_LOSS = 1000\n",

    "        OUT_TIME = 4\n",

    "        generator = loader\n",

    "        device_num = int(str(device)[-1])\n",

    "        dataset = loader._loader._loader.dataset\n",

    "    else:\n",

    "        tlen = len(loader)\n",

    "        OUT_LOSS = 50\n",

    "        OUT_TIME = 50\n",

    "        generator = enumerate(loader)\n",

    "        device_num = 1\n",

    "        dataset = loader.dataset\n",

    "    \n",

    "    #print('Start training {}'.format(device), 'batches', tlen)\n",

    "    \n",

    "    criterion = BCEWithLogitsLoss(weight = torch.Tensor(class_weights).to(device))\n",

    "    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(0.9, 0.99))\n",

    "    \n",

    "    model.train()\n",

    "    \n",

    "    if CLOUD and TPU:\n",

    "        tracker = xm.RateTracker()\n",

    "\n",

    "    tloss = 0\n",

    "    tloss_count = 0\n",

    "    \n",

    "    st = time.time()\n",

    "    mixup_collected = False\n",

    "    x_le_mix = None\n",

    "    lambd = None\n",

    "    for i, (x, x_le, y, non_black) in generator:\n",

    "        if (not CLOUD) or CLOUD_SINGLE:\n",

    "            x = x.to(device)\n",

    "            x_le = x_le.to(device)\n",

    "            y = y.to(device)\n",

    "            non_black = non_black.to(device)\n",

    "        \n",

    "        if MIXUP:\n",

    "            if mixup_collected:\n",

    "                lambd = np.random.beta(0.4, 0.4, y.size(0))\n",

    "                lambd = torch.Tensor(lambd).to(device)[:,None,None]\n",