{"cells":[{"metadata":{},"cell_type":"markdown","source":" I converted this notebook (https://www.kaggle.com/carlossouza/osic-autoencoder-training) to VAE. \n","execution_count":null},{"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.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# 2. Imports and global variables","execution_count":null},{"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 = 10\nval_size = 0.2\ntensorboard_dir = '/kaggle/working/runs'\n\n# Device configuration\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 3. Dataset interface","execution_count":null},{"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.**","execution_count":null},{"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","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, 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).","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 {'image': image, 'metadata': data}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 3.2.3. resize.py","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, 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","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, 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.","execution_count":null},{"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","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, 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","execution_count":null},{"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.","execution_count":null},{"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","execution_count":null},{"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":{"trusted":true},"cell_type":"code","source":"class VarAutoEncoder(nn.Module):\n    def __init__(self, latent_features=latent_features):\n        super(VarAutoEncoder, 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        self.fc2 = nn.Linear(10 * 10, latent_features)\n        \n        # Decoder\n        self.fc3 = 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        \n        mu = self.fc1(x)\n        log_var = self.fc2(x)\n        \n        if return_partials:\n            \n            return mu, log_var, up3out_shape, i1, up2out_shape, i2, up1out_shape, i3, \\\n                   up0out_shape, i4\n\n        else:\n            return mu, log_var\n    \n    def reparameterize(self, mu, log_var):\n        std = torch.exp(log_var/2)\n        eps = torch.randn_like(std)\n        return mu + eps * std\n    \n    def forward(self, x):\n        mu, log_var, up3out_shape, i1, up2out_shape, i2, \\\n        up1out_shape, i3, up0out_shape, i4 = self.encode(x)\n        \n        z = self.reparameterize(mu, log_var)\n       \n        # Decoder\n        x = F.relu(self.fc3(z))\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, mu, log_var","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.","execution_count":null},{"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 = VarAutoEncoder(latent_features=latent_features).to(device)\n\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, mu, log_var = model(inputs)\n                \n                # For KL divergence, see Appendix B in VAE paper or http://yunjey47.tistory.com/43\n                reconst_loss = F.mse_loss(outputs, inputs, size_average=False)\n                kl_div = - 0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())\n                \n                loss =  reconst_loss + kl_div\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.","execution_count":null},{"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    mu, log_var = model.encode(img, return_partials=False)\n    outputs, mu, log_var = model(img)\n    outputs = outputs.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\nHTML('<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","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}