{"cells":[{"metadata":{},"cell_type":"markdown","source":"# 1. OSIC AutoEncoder training\nThis notebooks demonstrates how to train a convolutional AutoEncoder to learn latent features from the 3D CT scans dataset.\n\nOne of the main applications of AutoEncoders is dimensionality reduction. We will use them for that: reducing 3D images (preprocessed to 1 x 40 x 256 x 256 tensors) to vectors (with 10 dimensions).\n![autoencoder](https://hackernoon.com/hn-images/1*8ixTe1VHLsmKB3AquWdxpQ.png)\n\nOnce we have the trained model, the idea is to apply it to extract these latent features and combine them with the OSIC tabular data.\n\nMy first experiments had a less strangled bottleneck (started with 96 x 2 x 20 x 20), which was already a reduction of over 34:1 (the inputs are 3D images of 1 x 40 x 256 x 256). The AutoEncoder output was great, easy to see. However, using latent features of 96 x 2 x 20 x 20 meant that, in the tabular model, I had to combine 76,800 features (flattened) with the 9 tabular features. In order to have a better balance between tabular and latent features, I decide to strangle the bottleneck further, squeezing the 3D images to 10 features (already flatenned in the AutoEncoder model). As you can see below, the model learns as the loss keeps going down. However, the output of the AutoEncoder is not as visible as with the less strangled bottleneck."},{"metadata":{},"cell_type":"markdown","source":"# 2. Imports and global variables"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import copy\nfrom datetime import timedelta, datetime\nimport imageio\nimport matplotlib.pyplot as plt\nfrom matplotlib import cm\nimport multiprocessing\nimport numpy as np\nimport os\nfrom pathlib import Path\nimport pydicom\nimport pytest\nimport scipy.ndimage as ndimage\nfrom scipy.ndimage.interpolation import zoom\nfrom skimage import measure, morphology, segmentation\nfrom time import time, sleep\nfrom tqdm import trange, tqdm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, random_split, DistributedSampler, DataLoader\nfrom torch.utils.tensorboard import SummaryWriter\nfrom torchvision import transforms\nimport warnings","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"root_dir = '/kaggle/input/osic-cached-dataset'\ntest_dir = '/kaggle/input/osic-pulmonary-fibrosis-progression/test'\nmodel_file = '/kaggle/working/diophantus.pt'\nresize_dims = (40, 256, 256)\nclip_bounds = (-1000, 200)\nwatershed_iterations = 1\npre_calculated_mean = 0.02865046213070556\nlatent_features = 10\nbatch_size = 16\nlearning_rate = 3e-5\nnum_epochs = 100\nval_size = 0.2\ntensorboard_dir = '/kaggle/working/runs'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 3. Dataset interface"},{"metadata":{},"cell_type":"markdown","source":"## 3.1. ctscans_dataset.py\nThis interface ingests the data from the 3D CT scans, porting them to a PyTorch Dataset.**"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CTScansDataset(Dataset):\n    def __init__(self, root_dir, transform=None):\n        self.root_dir = Path(root_dir)\n        self.patients = [p for p in self.root_dir.glob('*') if p.is_dir()]\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.patients)\n\n    def __getitem__(self, idx):\n        if torch.is_tensor(idx):\n            idx = idx.tolist()\n\n        image, metadata = self.load_scan(self.patients[idx])\n        sample = {'image': image, 'metadata': metadata}\n        if self.transform:\n            sample = self.transform(sample)\n\n        return sample\n\n    def save(self, path):\n        t0 = time()\n        Path(path).mkdir(exist_ok=True, parents=True)\n        print('Saving pre-processed dataset to disk')\n        sleep(1)\n        cum = 0\n\n        bar = trange(len(self))\n        for i in bar:\n            sample = self[i]\n            image, data = sample['image'], sample['metadata']\n            cum += torch.mean(image).item()\n\n            bar.set_description(f'Saving CT scan {data.PatientID}')\n            fname = Path(path) / f'{data.PatientID}.pt'\n            torch.save(image, fname)\n\n        sleep(1)\n        bar.close()\n        print(f'Done! Time {timedelta(seconds=time() - t0)}\\n'\n              f'Mean value: {cum / len(self)}')\n\n    def get_patient(self, patient_id):\n        patient_ids = [str(p.stem) for p in self.patients]\n        return self.__getitem__(patient_ids.index(patient_id))\n\n    @staticmethod\n    def 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.2. Pre-processing\nThere are some pre-processing to be done. Let's tackle them one step at a time.\n### 3.2.1. crop_bounding_box.py"},{"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, data = sample['image'], sample['metadata']\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 {'image': image, 'metadata': data}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.2.2. convert_to_hu.py\nCredits to [Guido Zuidhof's tutorial](https://www.kaggle.com/gzuidhof/full-preprocessing-tutorial)."},{"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 {'image': image, 'metadata': data}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.2.3. resize.py"},{"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, data = sample['image'], sample['metadata']\n        resize_factor = np.array(self.output_size) / np.array(image.shape)\n        image = zoom(image, resize_factor, mode='nearest')\n        return {'image': image, 'metadata': data}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.2.4. clip.py"},{"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, data = sample['image'], sample['metadata']\n        image[image < self.min] = self.min\n        image[image > self.max] = self.max\n        return {'image': image, 'metadata': data}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.2.5. mask_watershed.py\nCredits to [Aadhav Vignesh's amazing kernel](https://www.kaggle.com/aadhavvignesh/lung-segmentation-by-marker-controlled-watershed).\n\nIMPORTANT: I made some changes in Vignesh's code below to make it scalable, most notably reducing the number of iterations from 8 to 1. This was important to reduce the time to generate masks from ~8-9 seconds/slice (which would take over 17 hours to complete) to ~100ms/slice. I'm satisfied with the quality of the masks, as you can see in some samples below. However, using 8 iterations generate even better masks."},{"metadata":{"trusted":true},"cell_type":"code","source":"class MaskWatershed:\n    def __init__(self, min_hu, iterations, show_tqdm):\n        self.min_hu = min_hu\n        self.iterations = iterations\n        self.show_tqdm = show_tqdm\n\n    def __call__(self, sample):\n        image, data = sample['image'], sample['metadata']\n\n        stack = []\n        if self.show_tqdm:\n            bar = trange(image.shape[0])\n            bar.set_description(f'Masking CT scan {data.PatientID}')\n        else:\n            bar = range(image.shape[0])\n        for slice_idx in bar:\n            sliced = image[slice_idx]\n            stack.append(self.seperate_lungs(sliced, self.min_hu,\n                                             self.iterations))\n\n        return {\n            'image': np.stack(stack),\n            'metadata': sample['metadata']\n        }\n\n    @staticmethod\n    def seperate_lungs(image, min_hu, iterations):\n        h, w = image.shape[0], image.shape[1]\n\n        marker_internal, marker_external, marker_watershed = MaskWatershed.generate_markers(image)\n\n        # Sobel-Gradient\n        sobel_filtered_dx = ndimage.sobel(image, 1)\n        sobel_filtered_dy = ndimage.sobel(image, 0)\n        sobel_gradient = np.hypot(sobel_filtered_dx, sobel_filtered_dy)\n        sobel_gradient *= 255.0 / np.max(sobel_gradient)\n\n        watershed = morphology.watershed(sobel_gradient, marker_watershed)\n\n        outline = ndimage.morphological_gradient(watershed, size=(3,3))\n        outline = outline.astype(bool)\n\n        # Structuring element used for the filter\n        blackhat_struct = [[0, 0, 1, 1, 1, 0, 0],\n                           [0, 1, 1, 1, 1, 1, 0],\n                           [1, 1, 1, 1, 1, 1, 1],\n                           [1, 1, 1, 1, 1, 1, 1],\n                           [1, 1, 1, 1, 1, 1, 1],\n                           [0, 1, 1, 1, 1, 1, 0],\n                           [0, 0, 1, 1, 1, 0, 0]]\n\n        blackhat_struct = ndimage.iterate_structure(blackhat_struct, iterations)\n\n        # Perform Black Top-hat filter\n        outline += ndimage.black_tophat(outline, structure=blackhat_struct)\n\n        lungfilter = np.bitwise_or(marker_internal, outline)\n        lungfilter = ndimage.morphology.binary_closing(lungfilter, structure=np.ones((5,5)), iterations=3)\n\n        segmented = np.where(lungfilter == 1, image, min_hu * np.ones((h, w)))\n\n        return segmented  #, lungfilter, outline, watershed, sobel_gradient\n\n    @staticmethod\n    def generate_markers(image, threshold=-400):\n        h, w = image.shape[0], image.shape[1]\n\n        marker_internal = image < threshold\n        marker_internal = segmentation.clear_border(marker_internal)\n        marker_internal_labels = measure.label(marker_internal)\n\n        areas = [r.area for r in measure.regionprops(marker_internal_labels)]\n        areas.sort()\n\n        if len(areas) > 2:\n            for region in measure.regionprops(marker_internal_labels):\n                if region.area < areas[-2]:\n                    for coordinates in region.coords:\n                        marker_internal_labels[coordinates[0], coordinates[1]] = 0\n\n        marker_internal = marker_internal_labels > 0\n\n        # Creation of the External Marker\n        external_a = ndimage.binary_dilation(marker_internal, iterations=10)\n        external_b = ndimage.binary_dilation(marker_internal, iterations=55)\n        marker_external = external_b ^ external_a\n\n        # Creation of the Watershed Marker\n        marker_watershed = np.zeros((h, w), dtype=np.int)\n        marker_watershed += marker_internal * 255\n        marker_watershed += marker_external * 128\n\n        return marker_internal, marker_external, marker_watershed","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.2.6. normalize.py, to_tensor.py, zero_center.py"},{"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, data = sample['image'], sample['metadata']\n        image = image.astype(np.float)\n        image = (image - self.min) / (self.max - self.min)\n        return {'image': image, 'metadata': data}\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, data = sample['image'], sample['metadata']\n        if self.add_channel:\n            image = np.expand_dims(image, axis=0)\n\n        return {'image': torch.from_numpy(image), 'metadata': data}\n    \n    \nclass ZeroCenter:\n    def __init__(self, pre_calculated_mean):\n        self.pre_calculated_mean = pre_calculated_mean\n\n    def __call__(self, tensor):\n        return tensor - self.pre_calculated_mean","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.2.7. Inspecting some slices"},{"metadata":{"trusted":true},"cell_type":"code","source":"def show(list_imgs, cmap=cm.bone):\n    list_slices = []\n    for img3d in list_imgs:\n        slc = int(img3d.shape[0] / 2)\n        img = img3d[slc]\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        axs[i].axis('off')\n        \n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test = CTScansDataset(\n    root_dir=test_dir,\n    transform=transforms.Compose([\n        CropBoundingBox(),\n        ConvertToHU(),\n        Resize(resize_dims),\n        Clip(bounds=clip_bounds),\n        MaskWatershed(min_hu=min(clip_bounds), iterations=1, show_tqdm=True),\n        Normalize(bounds=clip_bounds)\n    ]))\n\nlist_imgs = [test[i]['image'] for i in range(len(test))]\nshow(list_imgs)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 3.3. Caching pre-processed images in the disk\nPre-processing all 176 3D CT scans take some time. Depending on the parameters we choose, it can take hours. \n\nWith the current choice of parameters, it takes around 15 minutes. To accelerate experimentation, I already pre-cached the images with the preprocessing parameters in this notebook, saving them in a [public dataset](https://www.kaggle.com/carlossouza/osic-cached-dataset). \n\nThis way, you can preprocess only once, and experiment with the same preprocessed tensors. The code to preprocess and cache images in the disk is:\n```\ndata = CTScansDataset(\n    root_dir=root_dir,\n    transform=transforms.Compose([\n        CropBoundingBox(),\n        ConvertToHU(),\n        Resize(size),\n        Clip(bounds=clip_bounds),\n        MaskWatershed(\n            min_hu=min(clip_bounds),\n            iterations=watershed_iterations,\n            show_tqdm=False),\n        Normalize(bounds=clip_bounds),\n        ToTensor()\n    ]))\ndata.save(dest_dir)\n```\n\nFrom this point on, we use the `CTTensorsDataset` as the interface to ingest the preprocessed tensors, taking the data to training."},{"metadata":{"trusted":true},"cell_type":"code","source":"class CTTensorsDataset(Dataset):\n    def __init__(self, root_dir, transform=None):\n        self.root_dir = Path(root_dir)\n        self.tensor_files = sorted([f for f in self.root_dir.glob('*.pt')])\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.tensor_files)\n\n    def __getitem__(self, item):\n        if torch.is_tensor(item):\n            item = item.tolist()\n\n        image = torch.load(self.tensor_files[item])\n        if self.transform:\n            image = self.transform(image)\n\n        return {\n            'patient_id': self.tensor_files[item].stem,\n            'image': image\n        }\n\n    def mean(self):\n        cum = 0\n        for i in range(len(self)):\n            sample = self[i]['image']\n            cum += torch.mean(sample).item()\n\n        return cum / len(self)\n\n    def random_split(self, val_size: float):\n        num_val = int(val_size * len(self))\n        num_train = len(self) - num_val\n        return random_split(self, [num_train, num_val])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.3.1. Checking data pipeline"},{"metadata":{"trusted":true},"cell_type":"code","source":"train = CTTensorsDataset(\n    root_dir=root_dir,\n    transform=ZeroCenter(pre_calculated_mean=pre_calculated_mean)\n)\ncum = 0\nfor i in range(len(train)):\n    sample = train[i]['image']\n    cum += torch.mean(sample).item()\n\nassert cum / len(train) == pytest.approx(0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 4. AutoEncoder\nCredits to [Srinjay Paul's great tutorial](https://srinjaypaul.github.io/3D_Convolutional_autoencoder_for_brain_volumes/), and lots of papers (I will link them later).\n\nAs mentioned, I strangled the bottleneck to force very few latent features (10). The image below shows the transformations:\n![autoencoder](https://i.ibb.co/2hYZFc1/autoencoder.jpg)"},{"metadata":{"trusted":true},"cell_type":"code","source":"class AutoEncoder(nn.Module):\n    def __init__(self, latent_features=latent_features):\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.conv4 = nn.Conv3d(96, 1, 1)\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        self.pool4 = nn.MaxPool3d(kernel_size=2, stride=2, return_indices=True)\n        self.fc1 = nn.Linear(10 * 10, latent_features)\n        # Decoder\n        self.fc2 = nn.Linear(latent_features, 10 * 10)\n        self.deconv0 = nn.ConvTranspose3d(1, 96, 1)\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.unpool0 = nn.MaxUnpool3d(kernel_size=2, stride=2)\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\n        x = self.conv2(x)\n        up2out_shape = x.shape\n        x, i2 = self.pool2(x)\n\n        x = self.conv3(x)\n        up1out_shape = x.shape\n        x, i3 = self.pool3(x)\n\n        x = self.conv4(x)\n        up0out_shape = x.shape\n        x, i4 = self.pool4(x)\n\n        x = x.view(-1, 10 * 10)\n        x = F.relu(self.fc1(x))\n\n        if return_partials:\n            return x, up3out_shape, i1, up2out_shape, i2, up1out_shape, i3, \\\n                   up0out_shape, i4\n\n        else:\n            return x\n\n    def forward(self, x):\n        x, up3out_shape, i1, up2out_shape, i2, \\\n        up1out_shape, i3, up0out_shape, i4 = self.encode(x)\n\n        # Decoder\n        x = F.relu(self.fc2(x))\n        x = x.view(-1, 1, 1, 10, 10)\n        x = self.unpool0(x, output_size=up0out_shape, indices=i4)\n        x = self.deconv0(x)\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":"# 5. Training\nI decided to take this opportunity to learn how to use TPUs. However, after days of intense frustration, I gave up. XLA documentation is very poor. Notebook examples are either so simple they are not useful at all, or so advanced/complex it is impossible to understand what is happening. The code frequently freezes, and it is impossible to know what is happening in the background…\n\nThe code below runs smoothly on GPU."},{"metadata":{},"cell_type":"markdown","source":"## 5.1. Monitoring on Tensorboard\nCredits to [Shivam Kumar tutorial](https://www.kaggle.com/shivam1600/tensorboard-on-kaggle)."},{"metadata":{"trusted":true},"cell_type":"code","source":"# !rm -rf ./logs/ \n# !mkdir ./logs/\n# # Download Ngrok to tunnel the tensorboard port to an external port\n# !wget https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip\n# !unzip ngrok-stable-linux-amd64.zip\n# pool = multiprocessing.Pool(processes = 10)\n# results_of_processes = [\n#     pool.apply_async(os.system, args=(cmd, ), callback=None) for cmd in [\n#         f\"tensorboard --logdir {tensorboard_dir}/ --host 0.0.0.0 --port 6006 &\",\n#         \"./ngrok http 6006 &\"\n#     ]\n# ]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ! curl -s http://localhost:4040/api/tunnels | python3 -c \\\n#     \"import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 5.2. Training loop\nIMPORTANT: For the sake of the demonstration, I'm running this training only for 10 epochs. To have usable results, we need at least 100 epochs."},{"metadata":{"trusted":true},"cell_type":"code","source":"t0 = time()\n\n# Load the data\ndata = CTTensorsDataset(\n    root_dir=root_dir,\n    transform=ZeroCenter(pre_calculated_mean=pre_calculated_mean)\n)\ntrain_set, val_set = data.random_split(val_size)\ndatasets = {'train': train_set, 'val': val_set}\ndataloaders = {\n    x: DataLoader(\n        datasets[x],\n        batch_size=batch_size,\n        shuffle=(x == 'train'),\n        num_workers=2\n    ) for x in ['train', 'val']}\n\ndataset_sizes = {x: len(datasets[x]) for x in ['train', 'val']}\n\n# Prepare for training\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nmodel = AutoEncoder(latent_features=latent_features).to(device)\ncriterion = torch.nn.MSELoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\nbest_model_wts = None\nbest_loss = np.inf\n\ndate_time = datetime.now().strftime(\"%Y%m%d-%H%M\")\nlog_dir = Path(tensorboard_dir) / f'{date_time}'\nwriter = SummaryWriter(log_dir)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Training loop\nfor epoch in range(num_epochs):\n\n    # Each epoch has a training and validation phase\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_preds = 0\n\n        # Iterate over data.\n        bar = tqdm(dataloaders[phase])\n        for inputs in bar:\n            bar.set_description(f'Epoch {epoch + 1} {phase}'.ljust(20))\n            inputs = inputs['image'].to(device, dtype=torch.float)\n\n            # zero the parameter gradients\n            optimizer.zero_grad()\n\n            # forward\n            # track history if only in train\n            with torch.set_grad_enabled(phase == 'train'):\n                outputs = model(inputs)\n                loss = criterion(outputs, inputs)\n\n                # backward + optimize only if in training phase\n                if phase == 'train':\n                    loss.backward()\n                    optimizer.step()\n\n            # statistics\n            running_loss += loss.item() * inputs.size(0)\n            running_preds += inputs.size(0)\n            bar.set_postfix(loss=f'{running_loss / running_preds:0.6f}')\n\n        epoch_loss = running_loss / dataset_sizes[phase]\n        writer.add_scalar(f'Loss/{phase}', epoch_loss, epoch)\n\n        # deep copy the model\n        if phase == 'val' and epoch_loss < best_loss:\n            best_loss = epoch_loss\n            best_model_wts = copy.deepcopy(model.state_dict())\n            torch.save(best_model_wts, model_file)\n\n# load best model weights\nmodel.load_state_dict(best_model_wts)\n\nprint(f'Done! Time {timedelta(seconds=time() - t0)}')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 6. Inference and inspection\nThe code below inspects a random sample. As mentioned, the quality can be improved by increasing the number of latent features. However, that will become a problem later when we combine the latent features with the tabular features in the Quant Model."},{"metadata":{"trusted":true},"cell_type":"code","source":"slc = 0.5\nsample_id = np.random.randint(len(data))\nprint(f'Inspecting CT Scan {data[sample_id][\"patient_id\"]}')\n\nfig, axs = plt.subplots(1, 2, figsize=(10, 7))\n\nsample = data[sample_id]['image'].squeeze(0).numpy()\naxs[0].imshow(sample[int(40 * slc), :, :], cmap=cm.bone)\naxs[0].axis('off')\nimageio.mimsave(\"sample_input.gif\", sample, duration=0.0001)\n\nwith torch.no_grad():\n    img = data[sample_id]['image'].unsqueeze(0).float().to(device)\n    latent_features = model.encode(img, return_partials=False)\\\n        .squeeze().cpu().numpy().tolist()\n    outputs = model(img).squeeze().cpu().numpy()\n\naxs[1].imshow(outputs[int(40 * slc), :, :], cmap=cm.bone)\naxs[1].axis('off')\n\nimageio.mimsave(\"sample_output.gif\", outputs, duration=0.0001)\n\nrmse = ((sample - outputs)**2).mean()\nplt.show()\nprint(f'Latent features: {latent_features} \\nLoss: {rmse}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# from IPython.display import HTML\n# HTML('<br/><img src=\"https://i.ibb.co/gFxgRq6/sample-input.gif\" style=\"float: left; width: 30%; margin-right: 1%; margin-bottom: 0.5em;\">'\n#      '<img src=\"https://i.ibb.co/Jm57fWw/sample-output.gif\" style=\"float: left; width: 30%; margin-right: 1%; margin-bottom: 0.5em;\">'\n#      '<p style=\"clear: both;\">')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 7. Next steps\n- Train longer: 10 epochs is not enough to achieve good results\n- Use the latent features in the Quant Model, and check how much they improve the predictions\n- Investigate/debug why some of the latent features are zero"},{"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}