{"cells":[{"metadata":{},"cell_type":"markdown","source":""},{"metadata":{},"cell_type":"markdown","source":"# Simple PyTorch Pipeline
\nUsing all 8 cores of TPU
"},{"metadata":{},"cell_type":"markdown","source":"# Table of contents \n\n1. [Introduction](#1)\n2. [Import Packages](#2)\n3. [Loading Data](#3)\n4. [Model](#4)\n5. [Training on 8-cores of TPU](#5)\n7. [References](#7)"},{"metadata":{},"cell_type":"markdown","source":"**Version Notes**:-\n* Version 12 : Not using bfloat16, since it is giving loss in a whole number format. Also, not using internet in this version.\n* Version 16 : Code for training the model on 1-core of the TPU and saving it for the inference.\n* Version 19 : Code for training the model on all 8-cores of the TPU and saving it for the inference.\n* Version 23 : Solved a bug (in Version 22, we were normalizing the mask 😛 and that was a mistake) and also trained the model with 25 epochs. In this version, you will see that loss drops at faster rate as compared to previous version. 🙂 \n* Version 24 : Added augs, using resnet34 encoder, increased learning rate. \n* Version 25 : Using seresnext50 encoder and Cosine Annealing LR scheduler. Let's train for 150 epochs.\n* Version 31 : Added Group K-Fold validation, dice metric and curated comments/documentation."},{"metadata":{},"cell_type":"markdown","source":"**Note:- Inference notebook for this TPU Pytorch pipeline will be out soon, with a decent LB score.**\n\n*Update:\nExperimented with the private copy of this NB with 4-fold resnext50 Unet model, 256 img size, Cosine Annealing LR scheduler, 60 epochs per fold and got **0.840** public LB score.* "},{"metadata":{},"cell_type":"markdown","source":"# 1. Introduction\n\nThis notebook will show you how to use TPU with PyTorch in this competition.\n\nTraning with 1-core of the TPU is demonstrated in the [16th version of this NB](https://www.kaggle.com/joshi98kishan/hubmap-pytorch-tpu-segmentation?scriptVersionId=48001107). In this NB, **we will be training the model on all the 8-cores of the TPU.**\n\nThe TPU which is provided with the kaggle kernel is **TPU v3-8** (3 for 3rd generation and 8 for 8 cores).\n\nSo, there are total 8 cores in this TPU. Each of the cores can run computations independently. \n\nThe heavylisting of doing multicore training is done by PyTorch XLA's modules :\n`DistributedSampler` and `ParallelLoader`. "},{"metadata":{},"cell_type":"markdown","source":"### Installing PyTorch XLA\n\nIn order to use TPU with PyTorch, we have to install [PyTorch XLA](https://github.com/pytorch/xla/) library. \n\nWe can install this library online by running the below commands:\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`\n\nNote:- Turn on the TPU before going ahead, otherwise you will get the error, '*Missing XLA configuration*'."},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-input":false,"_kg_hide-output":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\n\n# Removing setup-script and wheel files, as we don't need them now.\n!rm pytorch-xla-env-setup.py\n!rm torch-nightly-cp37-cp37m-linux_x86_64.whl \n!rm torch_xla-nightly-cp37-cp37m-linux_x86_64.whl \n!rm torchvision-nightly-cp37-cp37m-linux_x86_64.whl","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We will be training the model with [U-Net](https://arxiv.org/pdf/1505.04597.pdf) architecture having the pretrained encoder (resnet, serexnet etc.). This model is implemented in PyTorch ([here](https://github.com/qubvel/segmentation_models.pytorch)) in the form of library called **segmentation_models_pytorch**. \n\nJust to show how to install it offline, we will be installing it from source, becasue we will be needing this library in inference kernel and there we won't be able to use internet. The below commands and the associated [dataset](https://www.kaggle.com/vineeth1999/segmentationmodelspytorch) is borrowed from this [notebook](https://www.kaggle.com/vineeth1999/hubmap-pytorch-efficientunet-offline)."},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!mkdir -p /tmp/pip/cache/\n!cp ../input/segmentationmodelspytorch/segmentation_models/efficientnet_pytorch-0.6.3.xyz /tmp/pip/cache/efficientnet_pytorch-0.6.3.tar.gz\n!cp ../input/segmentationmodelspytorch/segmentation_models/pretrainedmodels-0.7.4.xyz /tmp/pip/cache/pretrainedmodels-0.7.4.tar.gz\n!cp ../input/segmentationmodelspytorch/segmentation_models/segmentation-models-pytorch-0.1.2.xyz /tmp/pip/cache/segmentation_models_pytorch-0.1.2.tar.gz\n!cp ../input/segmentationmodelspytorch/segmentation_models/timm-0.1.20-py3-none-any.whl /tmp/pip/cache/\n!cp ../input/segmentationmodelspytorch/segmentation_models/timm-0.2.1-py3-none-any.whl /tmp/pip/cache/\n!pip install --no-index --find-links /tmp/pip/cache/ efficientnet-pytorch\n!pip install --no-index --find-links /tmp/pip/cache/ segmentation-models-pytorch","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 2. Import Packages\n[Table of contents](#0.1)"},{"metadata":{"trusted":true},"cell_type":"code","source":"import os, gc\nimport numpy as np \nimport pandas as pd \nfrom tqdm.notebook import tqdm\nimport tifffile as tiff\nimport matplotlib.pyplot as plt\nimport cv2\nimport random\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch import optim\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision\nfrom torchvision import transforms\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts\n\nfrom albumentations import *\n\nfrom sklearn.model_selection import GroupKFold\n\nfrom segmentation_models_pytorch.unet import Unet\nfrom segmentation_models_pytorch.encoders import get_preprocessing_fn\n\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nfrom torch.utils.data.distributed import DistributedSampler\nimport torch_xla.distributed.parallel_loader as pl \nimport torch_xla.distributed.xla_multiprocessing as xmp\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Before going further, let's have a quick look on what is **bfloat16 floating-point format**.\n\n> This format is a truncated (16-bit) version of the 32-bit IEEE 754 single-precision floating-point format (binary32) with the intent of accelerating machine learning and near-sensor computing.\n\n> Bfloat16 is used to reduce the storage requirements and increase the calculation speed of machine learning algorithms. \n\n[wiki](https://en.wikipedia.org/wiki/Bfloat16_floating-point_format)"},{"metadata":{"trusted":true},"cell_type":"code","source":"'''Note:- Check the version notes above.'''\n\n# *** We set this environment variable, in order to make TPU use bfloat16. ***\n# os.environ['XLA_USE_BF16']=\"1\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def set_all_seeds(seed):\n random.seed(seed)\n os.environ['PYTHONHASHSEED'] = str(seed)\n np.random.seed(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n torch.backends.cudnn.deterministic = True\n \nset_all_seeds(2020)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"# 3. Loading Data\n[Table of contents](#0.1)"},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_PATH = '/kaggle/input/hubmap-256x256'\nos.listdir(DATA_PATH)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We are using iafoss's [256x 256 data](https://www.kaggle.com/iafoss/hubmap-256x256)."},{"metadata":{"trusted":true},"cell_type":"code","source":"PATH_TRAIN = os.path.join(DATA_PATH, 'train')\nPATH_MASKS = os.path.join(DATA_PATH, 'masks')\n\nprint(f'No. of training images: {len(os.listdir(PATH_TRAIN))}')\nprint(f'No. of masks: {len(os.listdir(PATH_MASKS))}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class HuBMAPDataset(Dataset):\n def __init__(self, \n data_path, \n fnames, \n preprocess_input = None,\n transforms = None):\n self.data_path = data_path\n self.fnames = fnames\n self.preprocess_input = preprocess_input\n self.transforms = transforms\n \n def __len__(self):\n return len(self.fnames)\n \n def __getitem__(self, idx):\n img_path = os.path.join(self.data_path, 'train', self.fnames[idx])\n mask_path = os.path.join(self.data_path, 'masks', self.fnames[idx])\n \n img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)\n mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)\n \n \n if self.transforms:\n # Applying augmentations if any. \n sample = self.transforms(image = img, \n mask = mask)\n \n img, mask = sample['image'], sample['mask']\n \n if self.preprocess_input:\n # Normalizing the image with the given mean and\n # std corresponding to each channel.\n img = self.preprocess_input(image = img)['image']\n \n # PyTorch assumes images in channels-first format. \n # Hence, bringing the channel at the first place.\n img = img.transpose((2, 0, 1))\n \n img = torch.from_numpy(img)\n mask = torch.from_numpy(mask)\n \n return img, mask","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Images and its corresponding masks are saved with the same filename.\nfnames = np.array(os.listdir(PATH_TRAIN))\n\ngroups = [fname[:9] for fname in fnames]\n \ngroup_kfold = GroupKFold(n_splits = 4)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Let's also have look on preprocessing.**\n\nWe are using pretrained model trained on imagenet images, normalized with these *mean* and *standard deviation*. \n\nmean : [0.485, 0.456, 0.406]\n\nstd : [0.229, 0.224, 0.225]\n\nFirst the whole image is divided by 255. Then from each channel, the corresponding mean is subtracted.\nAnd lastly, each channel is divided by its corresponding standard deviation (std).\n\n\n**img = ( (img/255) - mean)/std**"},{"metadata":{"trusted":true},"cell_type":"code","source":"# You can select and experiment with other encoder from this list:\n# https://github.com/qubvel/segmentation_models.pytorch#encoders\n\nENCODER_NAME = 'se_resnext50_32x4d'\n\n\n\n'''\nSo, we have to do the same kind of preprocessing while doing transfer learning. \n\nThe 'get_preprocessing_fn()' function will return the preprocessing_fn corresponding to the \nmentioned model and then we are using it as Albumentation like transform by wrapping it with \na Lambda function.\n'''\npreprocessing_fn = Lambda(image = get_preprocessing_fn(encoder_name = ENCODER_NAME,\n pretrained = 'imagenet'))\n\n\n\n# https://www.kaggle.com/iafoss/hubmap-pytorch-fast-ai-starter\ntransforms = Compose([\n HorizontalFlip(),\n VerticalFlip(),\n RandomRotate90(),\n ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.2, rotate_limit=15, p=0.9, \n border_mode=cv2.BORDER_REFLECT),\n OneOf([\n OpticalDistortion(p=0.3),\n GridDistortion(p=.1),\n IAAPiecewiseAffine(p=0.3),\n ], p=0.3),\n OneOf([\n HueSaturationValue(10,15,10),\n CLAHE(clip_limit=2),\n RandomBrightnessContrast(), \n ], p=0.3),\n ], p=1.0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 4. Model\n[Table of contents](#0.1)"},{"metadata":{},"cell_type":"markdown","source":"### U-Net model"},{"metadata":{},"cell_type":"markdown","source":"U-Net architecutre\n\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"class HuBMAPModel(nn.Module):\n def __init__(self):\n super(HuBMAPModel, self).__init__()\n self.model = Unet(encoder_name = ENCODER_NAME, \n encoder_weights = 'imagenet',\n classes = 1,\n activation = None)\n \n \n def forward(self, images):\n img_masks = self.model(images)\n return img_masks","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"In this competition, we have dice coefficient as the evaluation metric.\n\n**More the dice coefficient, better the predictions.**\n\nDice coefficient is pretty much straight forward.\n\n\nFor the training, we have to define the loss function. We can define the dice loss function like this:\n\n**dice loss = (1 - dice_coefficient)**\n\n**Lesser the dice loss, bigger the dice coefficient.**"},{"metadata":{"trusted":true},"cell_type":"code","source":"#https://www.kaggle.com/bigironsphere/loss-function-library-keras-pytorch\nclass DiceLoss(nn.Module):\n def __init__(self, weight=None, size_average=True):\n super(DiceLoss, self).__init__()\n\n def forward(self, inputs, targets, smooth=1):\n \n #comment out if your model contains a sigmoid or equivalent activation layer\n inputs = F.sigmoid(inputs) \n \n #flatten label and prediction tensors\n inputs = inputs.view(-1)\n targets = targets.view(-1)\n \n intersection = (inputs * targets).sum() \n dice = (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth) \n \n return 1 - dice\n \n \nloss_fn = DiceLoss()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_dice_coeff(pred, targs):\n '''\n Calculates the dice coeff of a single or batch of predicted mask and true masks.\n \n Args:\n pred : Batch of Predicted masks (b, w, h) or single predicted mask (w, h)\n targs : Batch of true masks (b, w, h) or single true mask (w, h)\n \n Returns: Dice coeff over a batch or over a single pair.\n '''\n \n \n pred = (pred>0).float()\n return 2.0 * (pred*targs).sum() / ((pred+targs).sum() + 1.0)\n\n \ndef reduce(values):\n ''' \n Returns the average of the values.\n Args:\n values : list of any value which is calulated on each core \n '''\n return sum(values) / len(values)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Encoder pretrained weights takes times to download, therefore taking the weigths from kaggle dataset\n\n!mkdir -p /root/.cache/torch/hub/checkpoints/\n!cp ../input/pretrained-model-weights-pytorch/se_resnext50_32x4d-a260b3a4.pth /root/.cache/torch/hub/checkpoints/","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 5. Training on 8-cores of TPU\n[Table of contents](#0.1)"},{"metadata":{},"cell_type":"markdown","source":"Let's see how Cosine Annealing LR scheduler changes learning rate with the epochs."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Did some experiments and come up with this set of hyperparams.\n\ndemo_model = HuBMAPModel()\n\n# LR which we set here is the highest LR value.\ndemo_optim = optim.SGD(demo_model.parameters(),\n lr = 0.3,\n momentum = 0.9)\nDEMO_EPOCHS = 60\n\n# Change the T_0 argument to get more number of cycles.\ndemo_sched = CosineAnnealingWarmRestarts(demo_optim, \n T_0 = DEMO_EPOCHS//3, \n T_mult=1, \n eta_min=0, \n last_epoch=-1, \n verbose=False)\nlrs = []\n\nfor i in range(DEMO_EPOCHS):\n demo_optim.step()\n lrs.append(demo_optim.param_groups[0][\"lr\"])\n demo_sched.step()\n\nplt.plot(lrs)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Updated\ndef train_one_epoch(epoch_no, data_loader, model, optimizer, device, scheduler = None):\n '''\n Run one epoch on the 'model'.\n Args:\n epoch_no: Serial number of the given epoch\n data_loader: Data iterator like DataLoader\n model: Model which needs to be train for one epoch\n optmizer : Pytorch's optimizer\n device: Device(A particular core) on which to run this epoch\n Scheduler : Pytorch's lr scheduler\n \n \n Returns: Nothing\n '''\n model.train()\n losses = []\n dice_coeffs = []\n \n for i, batch in enumerate(data_loader):\n img_btch, mask_btch = batch\n \n # *** Putting the images and masks to the TPU device. ***\n img_btch = img_btch.to(device)\n mask_btch = mask_btch.to(device)\n \n optimizer.zero_grad()\n \n pred_mask_btch = model(img_btch.float())\n \n loss = loss_fn(pred_mask_btch, mask_btch.float())\n \n loss.backward()\n \n \n '''\n xm.optimizer_step():\n Consolidates the gradients between cores and issues the XLA device step computation.\n The `step()` function now not only propagates gradients, but uses the TPU context \n to synchronize gradient updates across each processes' copy of the network. \n This ensures that each processes' network copy stays \"in sync\" (they are all identical).\n This means that each process's network has the same weights after this is called.\n [Source:PyTorch XLA doc]\n '''\n # Note: barrier=True not needed when using ParallelLoader\n xm.optimizer_step(optimizer)\n if scheduler is not None:\n scheduler.step()\n\n \n #'mesh_reduce()' reduce the loss calculated on 8 cores.\n # The way it needs to be reduced is defined in 'reduce()' function\n loss_reduced = xm.mesh_reduce('train_loss_reduce', \n loss, \n reduce)\n losses.append(loss_reduced.item())\n \n dice_coeff = get_dice_coeff(torch.squeeze(pred_mask_btch), \n mask_btch.float())\n dice_coeffs.append(xm.mesh_reduce('train_dice_reduce', \n dice_coeff, \n reduce).item()) \n \n \n del img_btch, pred_mask_btch, mask_btch\n gc.collect()\n \n xm.master_print(f'{epoch_no+1} - Loss : {reduce(losses): .4f}, Dice Coeff : {reduce(dice_coeffs): .4f}')\n \n\n# New stuff\ndef eval_fn(data_loader, model, device):\n '''\n Calculates metrics on the validation data.\n \n Returns: returns calculated metrics\n '''\n model.eval()\n \n dice_coeffs = []\n losses = []\n \n for i, batch in enumerate(data_loader):\n img_btch, mask_btch = batch\n \n img_btch = img_btch.to(device)\n mask_btch = mask_btch.to(device)\n \n pred_mask_btch = model(img_btch.float())\n \n loss = loss_fn(pred_mask_btch, \n mask_btch.float())\n losses.append(xm.mesh_reduce('val_loss_reduce', \n loss, \n reduce).item())\n \n dice_coeff = get_dice_coeff(torch.squeeze(pred_mask_btch), \n mask_btch.float())\n dice_coeffs.append(xm.mesh_reduce('val_dice_reduce', \n dice_coeff, \n reduce).item())\n \n total_dice_coeff = reduce(dice_coeffs)\n total_loss = reduce(losses)\n \n return total_loss, total_dice_coeff","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Note : _mp_fn() function will use xm.master_print() method to print any message, in place of print(). \n Because this function will run on each core, so print() would be called 8 times and print \n the message for each 8 core. xm.master_print() will print only for one time with respect to the\n master device."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Updated\ndef _mp_fn(rank, flags):\n \n ''' \n This function will be called for each device which takes part of the replication. \n \n Args:\n rank: Index of the process within the replication.\n flags: Contains custom arguments which you need to pass to each process.\n '''\n \n \n # Acquires the (unique) TPU core corresponding to this process's index\n device = xm.xla_device()\n \n # Creates the (distributed) train sampler, which let this process only access\n # its portion of the training dataset.\n train_sampler = DistributedSampler(dataset = flags['TRAIN_DS'],\n num_replicas = xm.xrt_world_size(),\n rank = xm.get_ordinal(),\n shuffle = True)\n train_dl = DataLoader(dataset = flags['TRAIN_DS'],\n batch_size = flags['BATCH_SIZE'],\n sampler = train_sampler,\n num_workers = 0)\n \n \n val_sampler = DistributedSampler(dataset = flags['VAL_DS'],\n num_replicas = xm.xrt_world_size(),\n rank = xm.get_ordinal(),\n shuffle = False)\n val_dl = DataLoader(dataset = flags['VAL_DS'],\n batch_size = flags['BATCH_SIZE'],\n sampler = val_sampler,\n num_workers = 0)\n \n del train_sampler, val_sampler\n gc.collect()\n \n# Uncomment this line, if you want to realize the training on all the 8-cores.\n# print(\"Process is using\", xm.xla_real_devices([str(device)])[0])\n \n fold_model = flags['FOLD_MODEL']\n fold_model.to(device)\n \n lr = flags['LR']\n\n optimizer = optim.SGD(model.parameters(), \n lr = lr,\n momentum = 0.9)\n \n scheduler = CosineAnnealingWarmRestarts(optimizer, \n T_0 = flags['EPOCHS']//3, \n T_mult=1, \n eta_min=0, \n last_epoch=-1, \n verbose=False)\n \n xm.master_print('Training now...')\n for e_no, epoch in enumerate(range(flags['EPOCHS'])):\n \n # Here comes our data loader for 8 cores.\n # It takes famous 'DataLoader()' object and list of \n # devices where data has to be sent.\n # Calling 'per_device_loader()' on it will\n # return the data loader for the particular device.\n train_para_loader = pl.ParallelLoader(train_dl, \n [device]).per_device_loader(device)\n \n train_one_epoch(e_no, \n train_para_loader,\n fold_model, \n optimizer, \n device,\n scheduler)\n \n del train_para_loader\n gc.collect()\n \n xm.master_print('\\nValidating now...')\n val_para_loader = pl.ParallelLoader(val_dl, \n [device]).per_device_loader(device)\n \n loss, dice_coeff = eval_fn(val_para_loader,\n fold_model,\n device)\n del val_para_loader\n gc.collect() \n \n xm.master_print(f'Val Loss : {loss: .4f}, Val Dice : {dice_coeff: .4f}')\n \n #Saving the model, so that we can import it in the inference kernel.\n xm.save(fold_model.state_dict(), f\"8core_fold_model_{flags['FOLD_NO']}.pth\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"✨✨✨\n\n**Now, we will be calling the most important function of this NB, which is `xmp.spawn()` for each fold.**\n\nThis function will create eight processes, one for each TPU core, and call above `_map_fn()` (map function) on each process. The inputs to `_map_fn()` are an index (zero through seven) and the placeholder `flags.` When the proccesses acquire their device they actually acquire their corresponding TPU core automatically. Each core is itself a device.\n\nAruguments of xmp.spawn():\n* fn = The function to be called for each device (TPU core).\n* args = The arguments which we will be passed to the `_mp_fn` function.\n* nprocs : The number of processes/devices for the replication.\n* start_method : The Python *multiprocessing* process creation method. For more info, visit [here](https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods)."},{"metadata":{"trusted":true},"cell_type":"code","source":"for fold, (t_idx, v_idx) in enumerate(group_kfold.split(fnames, \n groups = groups)):\n \n print(f'Fold: {fold+1}')\n print('-' * 40)\n \n t_fnames = fnames[t_idx]\n v_fnames = fnames[v_idx]\n \n train_ds = HuBMAPDataset(data_path = DATA_PATH,\n fnames = t_fnames,\n preprocess_input = preprocessing_fn,\n transforms = transforms)\n \n val_ds = HuBMAPDataset(data_path = DATA_PATH,\n fnames = v_fnames,\n preprocess_input = preprocessing_fn,\n transforms = None)\n \n model = HuBMAPModel()\n model.float()\n \n lr = 0.3\n \n FLAGS = {'FOLD_NO': fold,\n 'TRAIN_DS': train_ds,\n 'VAL_DS': val_ds,\n 'FOLD_MODEL': model,\n 'LR': lr,\n 'BATCH_SIZE' : 32,\n 'EPOCHS' : 30}\n\n xmp.spawn(fn = _mp_fn, \n args = (FLAGS,), \n nprocs = 8,\n start_method = 'fork')\n \n \n print('\\n')\n \n del train_ds, val_ds, model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# We have trained our warriors and they are resting in the current\n# working directory.\n# Now, we will move them to the battle ground (Inference kernel) for the war.\n!ls","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**So, this is how we can train our model on all the 8 cores of the TPU with PyTorch.**"},{"metadata":{},"cell_type":"markdown","source":"Do consider the upvote, if you liked it. :)"},{"metadata":{},"cell_type":"markdown","source":"# 7. References\n[Table of contents](#0.1)\n* https://www.kaggle.com/joshi98kishan/pytorch-xla-setup-script\n* https://www.kaggle.com/tanlikesmath/the-ultimate-pytorch-tpu-tutorial-jigsaw-xlm-r\n* https://pytorch.org/xla\n* https://www.kaggle.com/abhishek/super-duper-fast-pytorch-tpu-kernel\n* https://www.kaggle.com/abhishek/bert-multi-lingual-tpu-training-8-cores"}],"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}