{"cells":[{"metadata":{},"cell_type":"markdown","source":"# 1. End-to-end model: CT scans latent features + Tabular features\nThis model is a simple end-to-end solution that uses CT scans latent features and tabular features to generate predictions.\n\nThe CT scans latent features are obtained by encoding pre-processed 3D CT scan tensors, using a pre-trained AutoEncoder for that. \n\nThe tabular features are obtained by pre-processing the tabular data (more about it [here](https://www.kaggle.com/carlossouza/quantile-regression-pytorch-tabular-data-only)).\n\nI am not including here the code used to train the AutoEncoder. I pretrained it in another notebook, and saved the weights in a [public dataset](https://www.kaggle.com/carlossouza/osicautoencoder).\n\nThe overal model looks like the image below:\n![image](https://i.ibb.co/N65VVMp/final-5f18d00d86a6870013068112-270220.gif)\n\nThis version runs on CPU, end-to-end. However, it is slow. Improvements/next steps at the end.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# 2. Imports and global variables","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"use_TPU = False","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if use_TPU:\n    !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n    !python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if use_TPU:\n    import torch_xla\n    import torch_xla.debug.metrics as met\n    import torch_xla.distributed.parallel_loader as pl\n    import torch_xla.utils.utils as xu\n    import torch_xla.core.xla_model as xm\n    import torch_xla.distributed.xla_multiprocessing as xmp\n    import torch_xla.test.test_utils as test_utils\n\n    import warnings\n    warnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if use_TPU:\n    !export XLA_USE_BF16=1","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import copy\nfrom datetime import datetime, timedelta\nimport matplotlib.pyplot as plt\nfrom matplotlib import cm\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nimport pickle\nimport pydicom\nfrom sklearn.model_selection import GroupKFold, GroupShuffleSplit\nfrom torch.utils.data import Dataset\nfrom tqdm import tqdm\nimport torch\nfrom torch.utils.data import DataLoader, Subset\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom tqdm import trange\nfrom time import time\nimport warnings\nfrom scipy.ndimage.interpolation import zoom\nfrom enum import Enum\nfrom torchvision import transforms\nfrom skimage.measure import label, regionprops\nfrom skimage.segmentation import clear_border\nimport pytest","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"root_dir = Path('/kaggle/input/osic-pulmonary-fibrosis-progression')\nmodel_dir = Path('/kaggle/working')\npretrained_weigths_dir = Path('/kaggle/input/osicautoencoder')\npretrained_ae_weigths = pretrained_weigths_dir/'barcelona-20200722.pth'\ncache_dir = Path('/kaggle/input/osic-cached-dataset')\nlatent_dir = Path('/kaggle/working/latent')\nlatent_dir.mkdir(exist_ok=True, parents=True)\n# num_kfolds = 5\ntest_size=0.2\nbatch_size = 32\nlearning_rate = 1e-3\nnum_epochs = 20\nquantiles = (0.2, 0.5, 0.8)\nmodel_name ='cauchy'\n\nif use_TPU:\n    device = xm.xla_device()\nelse:\n    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 3. Dataset interface\nNow, the dataset interface not only ingests the tabular data, but also ingests and pre-process the 3D CT scans. The changes vs. [the baseline notebook](https://www.kaggle.com/carlossouza/quantile-regression-pytorch-tabular-data-only) are highlighted in the code.\n\nOn the code below, again credits to [Ulrich GOUE](https://www.kaggle.com/ulrich07).\n\nThe key trick here is to cache all tensors in disk, otherwise training will become extremely slow, as the bottleneck will become the DataLoader. That's because, without caching, every time ClinicalDataset is queried, it will re-preprocess and re-create 3D images from scratch. The solution below solves the problem. For perspective, without caching, loading a single sample takes ~1.5 seconds. With caching, ~6 millisencods.\n\nAnother solution would be just hold all 3D images on CPU memory, as they require approx 3.7 GB with the parameters I used in preprocessing pipeline, and transfer them to GPU/TPU when needed. May be even faster. I will experiment with that later.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class ClinicalDataset(Dataset):\n    def __init__(self, root_dir, ctscans_dir, mode, transform=None,\n                 cache_dir=None):\n        self.transform = transform\n        self.mode = mode\n        self.ctscans_dir = Path(ctscans_dir)\n        self.cache_dir = None if cache_dir is None else Path(cache_dir)\n\n        # If cache_dir is set, use cached values...\n        if cache_dir is not None:\n            self.raw = pd.read_csv(self.cache_dir/f'tabular_{mode}.csv')\n            with open(self.cache_dir/'features_list.pkl', \"rb\") as fp:\n                self.FE = pickle.load(fp)\n            return\n\n        # ...otherwise, pre-process\n        tr = pd.read_csv(Path(root_dir)/\"train.csv\")\n        tr.drop_duplicates(keep=False, inplace=True, subset=['Patient', 'Weeks'])\n        chunk = pd.read_csv(Path(root_dir)/\"test.csv\")\n\n        sub = pd.read_csv(Path(root_dir)/\"sample_submission.csv\")\n        sub['Patient'] = sub['Patient_Week'].apply(lambda x: x.split('_')[0])\n        sub['Weeks'] = sub['Patient_Week'].apply(lambda x: int(x.split('_')[-1]))\n        sub = sub[['Patient', 'Weeks', 'Confidence', 'Patient_Week']]\n        sub = sub.merge(chunk.drop('Weeks', axis=1), on=\"Patient\")\n\n        tr['WHERE'] = 'train'\n        chunk['WHERE'] = 'val'\n        sub['WHERE'] = 'test'\n        data = tr.append([chunk, sub])\n\n        data['min_week'] = data['Weeks']\n        data.loc[data.WHERE == 'test', 'min_week'] = np.nan\n        data['min_week'] = data.groupby('Patient')['min_week'].transform('min')\n\n        base = data.loc[data.Weeks == data.min_week]\n        base = base[['Patient', 'FVC']].copy()\n        base.columns = ['Patient', 'min_FVC']\n        base['nb'] = 1\n        base['nb'] = base.groupby('Patient')['nb'].transform('cumsum')\n        base = base[base.nb == 1]\n        base.drop('nb', axis=1, inplace=True)\n\n        data = data.merge(base, on='Patient', how='left')\n        data['base_week'] = data['Weeks'] - data['min_week']\n        del base\n\n        COLS = ['Sex', 'SmokingStatus']\n        self.FE = []\n        for col in COLS:\n            for mod in data[col].unique():\n                self.FE.append(mod)\n                data[mod] = (data[col] == mod).astype(int)\n\n        data['age'] = (data['Age'] - data['Age'].min()) / \\\n                      (data['Age'].max() - data['Age'].min())\n        data['BASE'] = (data['min_FVC'] - data['min_FVC'].min()) / \\\n                       (data['min_FVC'].max() - data['min_FVC'].min())\n        data['week'] = (data['base_week'] - data['base_week'].min()) / \\\n                       (data['base_week'].max() - data['base_week'].min())\n        data['percent'] = (data['Percent'] - data['Percent'].min()) / \\\n                          (data['Percent'].max() - data['Percent'].min())\n        self.FE += ['age', 'percent', 'week', 'BASE']\n\n        self.raw = data.loc[data.WHERE == mode].reset_index()\n        del data\n\n    def __len__(self):\n        return len(self.raw)\n\n    def __getitem__(self, idx):\n        if torch.is_tensor(idx):\n            idx = idx.tolist()\n\n        patient_id = self.raw['Patient'].iloc[idx]\n        if self.cache_dir is None:\n            patient_path = self.ctscans_dir / patient_id\n            image, metadata = load_scan(patient_path)\n        else:\n            image = torch.load(self.cache_dir / f'{patient_id}.pt')\n            metadata = pydicom.read_file(self.cache_dir / f'{patient_id}.dcm')\n\n        sample = {\n            'features': self.raw[self.FE].iloc[idx].values,\n            'image': image,\n            'metadata': metadata,\n            'target': self.raw['FVC'].iloc[idx]\n        }\n        if self.transform:\n            sample = self.transform(sample)\n\n        return sample\n\n    def cache(self, cache_dir):\n        Path(cache_dir).mkdir(exist_ok=True, parents=True)\n\n        # Cache raw features table\n        self.raw.to_csv(Path(cache_dir)/f'tabular_{self.mode}.csv', index=False)\n\n        # Cache features list\n        with open(Path(cache_dir)/'features_list.pkl', \"wb\") as fp:\n            pickle.dump(self.FE, fp)\n\n        # Cache images and metadata\n        self.raw['index'] = self.raw.index\n        idx_unique = self.raw.groupby('Patient').first()['index'].values\n        bar = tqdm(idx_unique.tolist())\n        for idx in bar:\n            sample = self[idx]\n            patient_id = sample['metadata'].PatientID\n            torch.save(sample['image'], Path(cache_dir)/f'{patient_id}.pt')\n            sample['metadata'].save_as(Path(cache_dir)/f'{patient_id}.dcm')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Helper function that loads CT scans in a single array. \n# This is also new vs. baselie\ndef load_scan(path):\n    slices = [pydicom.read_file(p) for p in path.glob('*.dcm')]\n    try:\n        slices.sort(key=lambda x: float(x.ImagePositionPatient[2]))\n    except AttributeError:\n        warnings.warn(f'Patient {slices[0].PatientID} CT scan does not '\n                      f'have \"ImagePositionPatient\". Assuming filenames '\n                      f'in the right scan order.')\n\n    image = np.stack([s.pixel_array.astype(float) for s in slices])\n    return image, slices[0]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 3.1. Preprocessing CT Scans\n\nThis section is new vs. the baseline notebook. It builds on a lot of great work done by several Kagglers (I'll make an effort and cite everyone).","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"### 3.1.1. Crop bounding boxes","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class CropBoundingBox:\n    @staticmethod\n    def bounding_box(img3d: np.array):\n        mid_img = img3d[int(img3d.shape[0] / 2)]\n        same_first_row = (mid_img[0, :] == mid_img[0, 0]).all()\n        same_first_col = (mid_img[:, 0] == mid_img[0, 0]).all()\n        if same_first_col and same_first_row:\n            return True\n        else:\n            return False\n\n    def __call__(self, sample):\n        image = sample['image']\n        if not self.bounding_box(image):\n            return sample\n\n        mid_img = image[int(image.shape[0] / 2)]\n        r_min, r_max = None, None\n        c_min, c_max = None, None\n        for row in range(mid_img.shape[0]):\n            if not (mid_img[row, :] == mid_img[0, 0]).all() and r_min is None:\n                r_min = row\n            if (mid_img[row, :] == mid_img[0, 0]).all() and r_max is None \\\n                    and r_min is not None:\n                r_max = row\n                break\n\n        for col in range(mid_img.shape[1]):\n            if not (mid_img[:, col] == mid_img[0, 0]).all() and c_min is None:\n                c_min = col\n            if (mid_img[:, col] == mid_img[0, 0]).all() and c_max is None \\\n                    and c_min is not None:\n                c_max = col\n                break\n\n        image = image[:, r_min:r_max, c_min:c_max]\n        return {\n            'features': sample['features'],\n            'image': image,\n            'metadata': sample['metadata'],\n            'target': sample['target']\n        }","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.1.2. Convert to Hounsfield Units\nGetting to the code below was really painful. Credits to lots of people.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class ConvertToHU:\n    def __call__(self, sample):\n        image, data = sample['image'], sample['metadata']\n\n        img_type = data.ImageType\n        is_hu = img_type[0] == 'ORIGINAL' and not (img_type[2] == 'LOCALIZER')\n        if not is_hu:\n            warnings.warn(f'Patient {data.PatientID} CT Scan not cannot be'\n                          f'converted to Hounsfield Units (HU).')\n\n        intercept = data.RescaleIntercept\n        slope = data.RescaleSlope\n        image = (image * slope + intercept).astype(np.int16)\n        return {\n            'features': sample['features'],\n            'image': image,\n            'metadata': data,\n            'target': sample['target']\n        }","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.1.3. Resize","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class Resize:\n    def __init__(self, output_size):\n        assert isinstance(output_size, tuple)\n        self.output_size = output_size\n\n    def __call__(self, sample):\n        image = sample['image']\n        resize_factor = np.array(self.output_size) / np.array(image.shape)\n        image = zoom(image, resize_factor, mode='nearest')\n        return {\n            'features': sample['features'],\n            'image': image,\n            'metadata': sample['metadata'],\n            'target': sample['target']\n        }","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.1.4. Clip","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class Clip:\n    def __init__(self, bounds=(-1000, 500)):\n        self.min = min(bounds)\n        self.max = max(bounds)\n\n    def __call__(self, sample):\n        image = sample['image']\n        image[image < self.min] = self.min\n        image[image > self.max] = self.max\n        return {\n            'features': sample['features'],\n            'image': image,\n            'metadata': sample['metadata'],\n            'target': sample['target']\n        }","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.1.5 Mask\nHere I basically tried 2 different methods: one [morphological](https://www.kaggle.com/miklgr500/unsupervise-lung-detection), and one using [Deep Learning](https://github.com/JoHof/lungmask). The deep learning method is extremely slow, so we go for the morphological path. The deep learning method is incomplete (thus the `raise NotImplementedError` in the code below). \n\nCredits to @[Michael Kazachok](https://www.kaggle.com/miklgr500).","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class MaskMethod(Enum):\n    MORPHOLOGICAL = 1\n    DEEPLEARNING = 2\n\n\nclass Mask:\n    def __init__(self, method=MaskMethod.MORPHOLOGICAL, threshold=-400,\n                 root_dir='../data/test'):\n        self.threshold = threshold\n        self.method = method\n        self.root_dir = root_dir\n\n    def __call__(self, sample):\n        image = sample['image']\n        if self.method == MaskMethod.MORPHOLOGICAL:\n            for slice_id in range(image.shape[0]):\n                m = self.get_morphological_mask(image[slice_id])\n                image[slice_id][m == False] = image[slice_id].min()\n        elif self.method == MaskMethod.DEEPLEARNING:\n            # m = self.get_deeplearning_mask(data.PatientID)\n            raise NotImplementedError\n        else:\n            raise ValueError('Unrecognized masking method')\n\n        return {\n            'features': sample['features'],\n            'image': image,\n            'metadata': sample['metadata'],\n            'target': sample['target']\n        }\n\n    def get_morphological_mask(self, image):\n        m = image < self.threshold\n        m = clear_border(m)\n        m = label(m)\n        areas = [r.area for r in regionprops(m)]\n        areas.sort()\n        if len(areas) > 2:\n            for region in regionprops(m):\n                if region.area < areas[-2]:\n                    for coordinates in region.coords:\n                        m[coordinates[0], coordinates[1]] = 0\n        return m > 0\n\n    def get_deeplearning_mask(self, patient_id):\n        \"\"\"Very slow, must be done using GPUs\n        \"\"\"\n        list_files = [str(x) for x in (Path(self.root_dir) / patient_id).glob('*.dcm')]\n        input_image = sitk.ReadImage(list_files)\n        m = mask.apply(input_image) #.squeeze()\n        m[m == 2] = 1\n        return m","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.1.6. Normalize, to tensor and zero center","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class Normalize:\n    def __init__(self, bounds=(-1000, 500)):\n        self.min = min(bounds)\n        self.max = max(bounds)\n\n    def __call__(self, sample):\n        image = sample['image'].astype(np.float)\n        image = (image - self.min) / (self.max - self.min)\n        return {\n            'features': sample['features'],\n            'image': image,\n            'metadata': sample['metadata'],\n            'target': sample['target']\n        }\n\n\nclass ToTensor:\n    def __init__(self, add_channel=True):\n        self.add_channel = add_channel\n\n    def __call__(self, sample):\n        image = sample['image']\n        if self.add_channel:\n            image = np.expand_dims(image, axis=0)\n\n        return {\n            'features': sample['features'],\n            'image': torch.from_numpy(image),\n            'metadata': sample['metadata'],\n            'target': sample['target']\n        }\n\n    \nclass ZeroCenter:\n    def __init__(self, pre_calculated_mean):\n        self.pre_calculated_mean = pre_calculated_mean\n\n    def __call__(self, sample):\n        return {\n            'features': sample['features'],\n            'image': sample['image'] - self.pre_calculated_mean,\n            'metadata': sample['metadata'],\n            'target': sample['target']\n        }","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 3.2. Checking the data pipeline\nBefore moving forward, let's ensure the data pipeline works","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def show(list_imgs, cmap=cm.bone, rgb=False):\n    list_slices = []\n    for img3d in list_imgs:\n        slc = int(img3d.shape[0] / 2)\n        img = img3d[slc]\n        if rgb:\n            img = (img * 255).astype(np.int16)\n        list_slices.append(img)\n    \n    fig, axs = plt.subplots(1, 5, figsize=(15, 7))\n    for i, img in enumerate(list_slices):\n        axs[i].imshow(img, cmap=cmap)\n    \n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data = ClinicalDataset(\n    root_dir=root_dir,\n    ctscans_dir=root_dir/'test',\n    mode='val',\n    transform=transforms.Compose([\n        CropBoundingBox(),\n        ConvertToHU(),\n        Resize((40, 256, 256)),\n        Clip(bounds=(-1000, 500)),\n        Mask(method=MaskMethod.MORPHOLOGICAL, threshold=-500),\n        Normalize(bounds=(-1000, -500))\n    ]))\n\nfor i in range(len(data)):\n    assert data[i]['image'].shape == (40, 256, 256)\n    \nlist_imgs = [data[i]['image'] for i in range(len(data))]\nshow(list_imgs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data = ClinicalDataset(\n    root_dir=root_dir,\n    ctscans_dir=root_dir/'test',\n    mode='val',\n    transform=transforms.Compose([\n        CropBoundingBox(),\n        ConvertToHU(),\n        Resize((40, 256, 256)),\n        Clip(bounds=(-1000, 500)),\n        Mask(method=MaskMethod.MORPHOLOGICAL, threshold=-500),\n        Normalize(bounds=(-1000, -500)),\n        ToTensor(),\n        ZeroCenter(pre_calculated_mean=0.029105728564346046)\n    ]))\n\nmeans = []\nfor i in range(len(data)):\n    means.append(data[i]['image'].mean())\n\nassert np.mean(means) == pytest.approx(0, abs=2e-3)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 4. Models and loss\nThe differences vs. the baseline notebook are highlighted in the code.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## 4.1. Quant model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class QuantModel(nn.Module):\n    def __init__(self, in_tabular_features=9, in_ctscan_features=76800,\n                 out_quantiles=3):\n        super(QuantModel, self).__init__()\n        # This line is new. We need to know a priori the number\n        # of latent features to properly flatten the tensor\n        self.in_ctscan_features = in_ctscan_features\n\n        self.fc1 = nn.Linear(in_tabular_features, 512)\n        self.fc2 = nn.Linear(in_ctscan_features, 512)\n        self.fc3 = nn.Linear(1024, 512)\n        self.fc4 = nn.Linear(512, out_quantiles)\n\n    def forward(self, x1, x2):\n        # Now the quant model has 2 inputs: x1 (the tabular features)\n        # and x2 (the pre-computed latent features)\n        x1 = F.relu(self.fc1(x1))\n        \n        # Flattens the latent features and concatenate with tabular features\n        x2 = x2.view(-1, self.in_ctscan_features)\n        x2 = F.relu(self.fc2(x2))\n        x = torch.cat([x1, x2], dim=1)\n        \n        x = F.relu(self.fc3(x))\n        x = self.fc4(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 4.2. Quant loss\nNo news here vs. baseline notebook. Again, credits to [Ceshine Lee's quantile regression tutorial](https://medium.com/the-artificial-impostor/quantile-regression-part-2-6fdbc26b2629).","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def quantile_loss(preds, target, quantiles):\n    assert not target.requires_grad\n    assert preds.size(0) == target.size(0)\n    losses = []\n    for i, q in enumerate(quantiles):\n        errors = target - preds[:, i]\n        losses.append(torch.max((q - 1) * errors, q * errors).unsqueeze(1))\n    loss = torch.mean(torch.sum(torch.cat(losses, dim=1), dim=1))\n    return loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 4.3. AutoEncoder\nThis is new vs. baseline notebook. As mentioned, the auto-encoder was pre-trained in another notebook. I won't include the pre-training code here otherwise this notebook will become even larger.\n\nCredits to [Srinjay Paul's 3D Convolutional autoencoder for brain volumes tutorial](https://srinjaypaul.github.io/3D_Convolutional_autoencoder_for_brain_volumes/).","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class AutoEncoder(nn.Module):\n    def __init__(self):\n        super(AutoEncoder, self).__init__()\n        # Encoder\n        self.conv1 = nn.Conv3d(1, 16, 3)\n        self.conv2 = nn.Conv3d(16, 32, 3)\n        self.conv3 = nn.Conv3d(32, 96, 2)\n        self.pool1 = nn.MaxPool3d(kernel_size=2, stride=2, return_indices=True)\n        self.pool2 = nn.MaxPool3d(kernel_size=3, stride=3, return_indices=True)\n        self.pool3 = nn.MaxPool3d(kernel_size=2, stride=2, return_indices=True)\n        # Decoder\n        self.deconv1 = nn.ConvTranspose3d(96, 32, 2)\n        self.deconv2 = nn.ConvTranspose3d(32, 16, 3)\n        self.deconv3 = nn.ConvTranspose3d(16, 1, 3)\n        self.unpool1 = nn.MaxUnpool3d(kernel_size=2, stride=2)\n        self.unpool2 = nn.MaxUnpool3d(kernel_size=3, stride=3)\n        self.unpool3 = nn.MaxUnpool3d(kernel_size=2, stride=2)\n\n    def encode(self, x, return_partials=True):\n        # Encoder\n        x = self.conv1(x)\n        up3out_shape = x.shape\n        x, i1 = self.pool1(x)\n        x = self.conv2(x)\n        up2out_shape = x.shape\n        x, i2 = self.pool2(x)\n        x = self.conv3(x)\n        up1out_shape = x.shape\n        x, i3 = self.pool3(x)\n\n        if return_partials:\n            return x, up3out_shape, i1, up2out_shape, i2, up1out_shape, i3\n        else:\n            return x\n\n    def forward(self, x):\n        x, up3out_shape, i1, up2out_shape, i2, up1out_shape, i3 = self.encode(x)\n\n        # Decoder\n        x = self.unpool1(x, output_size=up1out_shape, indices=i3)\n        x = self.deconv1(x)\n        x = self.unpool2(x, output_size=up2out_shape, indices=i2)\n        x = self.deconv2(x)\n        x = self.unpool3(x, output_size=up3out_shape, indices=i1)\n        x = self.deconv3(x)\n\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 4.4. Cache all pre-processed 3D CT Scans and pre-compute all latent features\nIt takes ~15 minutes to pre-process all 3D CT scan. To accelerate, I already pre-cached the dataset and made it available in [this public dataset](https://www.kaggle.com/carlossouza/osic-cached-dataset). If you want to do it, the code to cache the dataset is `dataset.cache(cache_dir)`. \n\nWith the dataset cached, let's pre-compute all latent features. (Not pre-computing them would lead to very slow training...):","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Helper function that generates all latent features\nclass GenerateLatentFeatures:\n    def __init__(self, autoencoder, latent_dir):\n        self.autoencoder = autoencoder\n        self.latent_dir = Path(latent_dir)\n        self.cache_dir = Path(cache_dir)\n\n    def __call__(self, sample):\n        patient_id = sample['metadata'].PatientID\n        cached_latent_file = self.latent_dir/f'{patient_id}_lat.pt'\n\n        if cached_latent_file.is_file():\n            latent_features = torch.load(cached_latent_file)\n        else:\n            with torch.no_grad():\n                img = sample['image'].float().unsqueeze(0)\n                latent_features = self.autoencoder.encode(\n                    img, return_partials=False).squeeze(0)\n            torch.save(latent_features, cached_latent_file)\n\n        return {\n            'tabular_features': sample['features'],\n            'latent_features': latent_features,\n            'target': sample['target']\n        }","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The code below loops through all samples in the dataset, pre-computes the latent features, and cache them in disk as well.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"autoencoder = AutoEncoder()\nautoencoder.load_state_dict(torch.load(\n    pretrained_ae_weigths,\n    map_location=torch.device('cpu')\n))\nautoencoder.to(device)\nautoencoder.eval()\n\ndata = ClinicalDataset(\n    root_dir=root_dir,\n    ctscans_dir=root_dir/'train',\n    cache_dir=cache_dir,\n    mode='train',\n    transform=GenerateLatentFeatures(autoencoder, latent_dir)\n)\nfor i in trange(len(data)):\n    sample = data[i]\n    assert sample['latent_features'].shape == (96, 2, 20, 20)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 4.5. Overfit a small batch before moving forward\nAs observed by Andrej Karpathy ([A Recipe for Training Neural Networks](http://karpathy.github.io/2019/04/25/recipe/)), neural net training fail silently. Before moving to full blown training, let's overfit a small batch: if we observe the loss decreasing, it is a sign that we probably did not make any obvious mistakes so far.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"dataloader = DataLoader(data, batch_size=batch_size,\n                        shuffle=True, num_workers=2)\nbatch = next(iter(dataloader))\n\nmodel = QuantModel().to(device)\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n\nbar = trange(50)\nfor epoch in bar:\n    inputs1 = batch['tabular_features'].float().to(device)\n    inputs2 = batch['latent_features'].float().to(device)\n    targets = batch['target'].to(device)\n\n    optimizer.zero_grad()\n    preds = model(inputs1, inputs2)\n    loss = quantile_loss(preds, targets, quantiles)\n    loss.backward()\n    if use_TPU:\n        xm.optimizer_step(optimizer, barrier=True)\n    else:\n        optimizer.step()\n    \n\n    bar.set_postfix(loss=f'{loss.item():0.1f}')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 5. Training\nI made 2 main changes here:\n1. Code is adjusted to run on TPU. It is not optimized (uses only 1 of the 8 cores), but does the job... Anyway, the last version does everything on CPU\n2. It is NOT doing Kfold CV. Code is back to doing simple Group Shuffle split. Main reason is that I did not see noticeable improvement from doing a simple split vs. KFold. When we were only using tabular data, no problem doing 5 folds. Now, every fold takes some time. As next step, an idea would be run every fold in parallel, in one of the 8 TPU cores, as I saw in one of the TPU tutorials. But that's for later","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Helper generator that group splits\ndef group_split(dataset, groups, test_size=0.2):\n    gss = GroupShuffleSplit(n_splits=1, test_size=test_size)\n    idx = list(gss.split(dataset.raw, dataset.raw, groups))\n    train = Subset(dataset, idx[0][0])\n    val = Subset(dataset, idx[0][1])\n    return train, val\n        \n# Helper function with competition metric\ndef metric(preds, targets):\n    sigma = preds[:, 2] - preds[:, 0]\n    sigma[sigma < 70] = 70\n    delta = (preds[:, 1] - targets).abs()\n    delta[delta > 1000] = 1000\n    return -np.sqrt(2) * delta / sigma - torch.log(np.sqrt(2) * sigma)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Load the data\nautoencoder = AutoEncoder()\nautoencoder.load_state_dict(torch.load(\n    pretrained_ae_weigths,\n    map_location=torch.device('cpu')\n))\nautoencoder.eval()\n\ndata = ClinicalDataset(\n    root_dir=root_dir,\n    ctscans_dir=root_dir/'train',\n    cache_dir=cache_dir,\n    mode='train',\n    transform=GenerateLatentFeatures(autoencoder, latent_dir)\n)\n\ntrainset, valset = group_split(data, data.raw['Patient'], test_size)\nt0 = time()\n\n# Prepare to save model weights\nPath(model_dir).mkdir(parents=True, exist_ok=True)\nnow = datetime.now()\nfname = f'{model_name}-{now.year}{now.month:02d}{now.day:02d}.pth'\nmodel_file = Path(model_dir) / fname\n\ndataset_sizes = {'train': len(trainset), 'val': len(valset)}\ndataloaders = {\n    'train': DataLoader(trainset, batch_size=batch_size,\n                        shuffle=True, num_workers=2),\n    'val': DataLoader(valset, batch_size=batch_size,\n                      shuffle=False, num_workers=2)\n}\n\n# Create the model and optimizer\nmodel = QuantModel().to(device)\noptimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n\n# Set global tracking variables\nepoch_loss = {'train': np.inf, 'val': np.inf}\nepoch_metric = {'train': -np.inf, 'val': -np.inf}\nbest_loss = np.inf\nbest_model_wts = None\ndf = pd.DataFrame(columns=['epoch', 'train_loss', 'val_loss'])\n\n# Training loop\nfor epoch in range(num_epochs):\n    for phase in ['train', 'val']:\n        if phase == 'train':\n            model.train()  # Set model to training mode\n        else:\n            model.eval()   # Set model to evaluate mode\n\n        running_loss = 0.0\n        running_metric = 0.0\n\n        # Iterate over data\n        num_samples = 0\n        bar = tqdm(dataloaders[phase])\n        for batch in bar:\n            bar.set_description(f'Epoch {epoch} {phase}'.ljust(20))\n            inputs1 = batch['tabular_features'].float().to(device)\n            inputs2 = batch['latent_features'].float().to(device)\n            targets = batch['target'].to(device)\n\n            # zero the parameter gradients\n            optimizer.zero_grad()\n            # forward\n            # track gradients if only in train\n            with torch.set_grad_enabled(phase == 'train'):\n                preds = model(inputs1, inputs2)\n                loss = quantile_loss(preds, targets, quantiles)\n                # backward + optimize only if in training phase\n                if phase == 'train':\n                    loss.backward()\n                    if use_TPU:\n                        xm.optimizer_step(optimizer, barrier=True)\n                    else:\n                        optimizer.step()\n\n            running_loss += loss.item() * inputs1.size(0)\n            running_metric += metric(preds, targets).sum()\n\n            # batch statistics\n            num_samples += inputs1.size(0)\n            bar.set_postfix(loss=f'{running_loss / num_samples:0.1f}',\n                            metric=f'{running_metric / num_samples:0.4f}')\n\n        # epoch statistics\n        epoch_loss[phase] = running_loss / dataset_sizes[phase]\n        epoch_metric[phase] = running_metric / dataset_sizes[phase]\n\n        # deep copy the model\n        if phase == 'val' and epoch_loss['val'] < best_loss:\n            best_loss = epoch_loss['val']\n            best_model_wts = copy.deepcopy(model.state_dict())\n            torch.save(best_model_wts, model_file)\n\n    df = df.append({\n        'epoch': epoch + 1,\n        'train_loss': epoch_loss[\"train\"],\n        'val_loss': epoch_loss[\"val\"]\n    }, ignore_index=True)\n\n# Save training statistics\nfname = f'{model_name}-{now.year}{now.month:02d}{now.day:02d}.csv'\ncsv_file = Path(model_dir) / fname\ndf.to_csv(csv_file, index=False)\n\n# load best model weights\nmodel.load_state_dict(best_model_wts)\n\nprint(f'Training complete! Time: {timedelta(seconds=time() - t0)}')\nmodels = [model]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 6. Generating submission CSV\nNo change vs. baseline notebook","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"data = ClinicalDataset(\n    root_dir=root_dir,\n    ctscans_dir=root_dir/'test',\n    mode='test',\n    transform=transforms.Compose([\n        CropBoundingBox(),\n        ConvertToHU(),\n        Resize((40, 256, 256)),\n        Clip(bounds=(-1000, 500)),\n        Mask(method=MaskMethod.MORPHOLOGICAL, threshold=-500),\n        Normalize(bounds=(-1000, -500)),\n        ToTensor(),\n        ZeroCenter(pre_calculated_mean=0.029105728564346046)\n    ]))\n\ndata.cache(latent_dir)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data = ClinicalDataset(\n    root_dir=root_dir,\n    ctscans_dir=root_dir/'test',\n    cache_dir=latent_dir,\n    mode='test',\n    transform=GenerateLatentFeatures(autoencoder, latent_dir)\n)\n\navg_preds = np.zeros((len(data), len(quantiles)))\n\nfor model in models:\n    dataloader = DataLoader(data, batch_size=batch_size,\n                            shuffle=False, num_workers=2)\n    preds = []\n    for batch in tqdm(dataloader):\n        inputs1 = batch['tabular_features'].float()\n        inputs2 = batch['latent_features'].float()\n        with torch.no_grad():\n            preds.append(model(inputs1, inputs2))\n\n    preds = torch.cat(preds, dim=0).numpy()\n    avg_preds += preds\n\navg_preds /= len(models)\ndf = pd.DataFrame(data=avg_preds, columns=list(quantiles))\ndf['Patient_Week'] = data.raw['Patient_Week']\ndf['FVC'] = df[quantiles[1]]\ndf['Confidence'] = df[quantiles[2]] - df[quantiles[0]]\ndf = df.drop(columns=list(quantiles))\ndf.to_csv('submission.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(len(df))\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 7. Conclusions and next steps\n- The process is working, end-to-end\n- However, it is very slow... need to figure out how to accelerate\n- Need to experiment/tune hyperparameters to improve metric\n- The final steps are very ugly, should refactor","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}