{"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.","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":"!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":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import torch_xla\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.utils.utils as xu\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.test.test_utils as test_utils\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\nimport os\nos.environ['XLA_USE_BF16']=\"1\"\nos.environ['XLA_TENSOR_ALLOCATOR_MAXSIZE'] = '100000000'","execution_count":null,"outputs":[]},{"metadata":{"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","execution_count":null,"outputs":[]},{"metadata":{"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 = 4\nlearning_rate = 3e-5\nnum_epochs = 10\nval_size = 0.2\ntensorboard_dir = '/kaggle/working/runs'\nflags = {\n    'batch_size': batch_size,\n    'num_epochs': num_epochs,\n    'seed': 1234,\n    'learning_rate': learning_rate,\n    'model_file': model_file\n}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 3. Dataset interface\n## 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":{},"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)","execution_count":null},{"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 on TPU\nTrying again... Credits to [ilovescience great tutorial](https://www.kaggle.com/tanlikesmath/the-ultimate-pytorch-tpu-tutorial-jigsaw-xlm-r).","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"mx = AutoEncoder(latent_features=10)\nloss_fn = torch.nn.MSELoss()\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)\n\n\ndef reduce_fn(vals):\n    return sum(vals) / len(vals)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_loop_fn(data_loader, model, optimizer, device):\n    model.train()\n    for batch_num, batch in enumerate(data_loader):\n\n        inputs = batch['image'].float().to(device)\n\n        # pass ids to model\n        optimizer.zero_grad()\n        outputs = model(inputs)\n        loss = loss_fn(outputs, inputs)\n\n        if batch_num % 20 == 0:\n            # since the loss is on all 8 cores, reduce the loss values\n            # and print the average (as defined in reduce_fn)\n            loss_reduced = xm.mesh_reduce('loss_reduce', loss, reduce_fn)\n            # master_print will only print once (not from all 8 cores)\n            xm.master_print(f'{batch_num}: loss={loss_reduced:0.6f}')\n\n        loss.backward()\n        xm.optimizer_step(optimizer)\n\n    model.eval() # put model in eval mode for later use","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def eval_loop_fn(data_loader, model, device):\n    with torch.no_grad():\n        loss_fn = torch.nn.MSELoss()\n        running_loss = 0.0\n        for batch_num, batch in enumerate(data_loader):\n            inputs = batch['image'].float().to(device)\n            outputs = model(inputs)\n            loss = loss_fn(outputs, inputs)\n            running_loss += loss * inputs.size(0)\n\n    return running_loss / len(data_loader.dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def run(index, flags):\n    # Sets a common random seed - both for initialization and\n    # ensuring graph is the same\n    torch.manual_seed(flags['seed'])\n\n    # Creates the (distributed) train sampler, which let this process\n    # only access its portion of the training dataset\n    train_sampler = DistributedSampler(\n        train_set,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True)\n\n    val_sampler = DistributedSampler(\n        val_set,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=False)\n\n    # Creates dataloaders, which load data in batches\n    # Note: test loader is not shuffled or sampled\n    train_loader = torch.utils.data.DataLoader(\n        train_set,\n        batch_size=flags['batch_size'],\n        sampler=train_sampler,\n        num_workers=0,\n        drop_last=True)  #\n\n    val_loader = torch.utils.data.DataLoader(\n        val_set,\n        batch_size=flags['batch_size'],\n        sampler=val_sampler,\n        num_workers=0,\n        drop_last=False)\n\n    # Acquires the (unique) Cloud TPU core corresponding\n    # to this process's index\n    device = xm.xla_device()\n\n    model = mx.to(device) # put model onto the TPU core\n    xm.master_print('done loading model')\n\n    lr = flags['learning_rate'] * xm.xrt_world_size()  # scale the learning rate\n    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n    \n    xm.master_print(met.metrics_report())\n\n    ## Trains\n    t0 = time()\n    for epoch in range(flags['num_epochs']):\n        para_loader = pl.ParallelLoader(train_loader, [device])\n        xm.master_print('parallel loader created... training now')\n        # call training loop:\n        train_loop_fn(para_loader.per_device_loader(device),\n                      model, optimizer, device)\n        \n        xm.master_print(met.metrics_report())\n\n        del para_loader\n        gc.collect()\n\n        para_loader = pl.ParallelLoader(val_loader, [device])\n        # call evaluation loop\n        loss = eval_loop_fn(para_loader.per_device_loader(device),\n                            model, device)\n\n        del para_loader\n        gc.collect()\n\n        # stats\n        loss_reduced = xm.mesh_reduce('loss_reduce', loss, reduce_fn)\n        xm.master_print(f'Epoch {epoch} loss: {loss_reduced:0.6f}')\n        gc.collect()\n\n    print(f\"Process {index} finished evaluation. Time: {time() - t0}\")\n\n    # save our model\n    xm.save(model.state_dict(), flags['model_file'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def map_fn(index, flags):\n    a = run(index, flags)\n    \n    \nxmp.spawn(map_fn, args=(flags,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"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}