{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"# installing the torch-xla nightly version\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 --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# making the necessary imports\nimport os\nimport torch\nimport pandas as pd\nfrom scipy import stats\nimport numpy as np\n\nfrom collections import OrderedDict, namedtuple\nimport torch.nn as nn\nfrom torch.optim import lr_scheduler\nimport joblib\n\nimport logging\nimport transformers\nfrom transformers import AdamW, get_linear_schedule_with_warmup, get_constant_schedule\nimport sys\nfrom sklearn import metrics, model_selection\n\nimport warnings\nimport torch_xla\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.data_parallel as dp\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\nimport warnings\n\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# separating out the filenames of training, validation and dataset separately in 3 different list variables\n\nimport glob\n\ntrain_files = glob.glob('/kaggle/input/tpu-getting-started/*/train/*.tfrec')\nval_files = glob.glob('/kaggle/input/tpu-getting-started/*/val/*.tfrec')\ntest_files = glob.glob('/kaggle/input/tpu-getting-started/*/test/*.tfrec')","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import tensorflow as tf\n\n# Create a dictionary describing the features.\ntrain_feature_description = {\n    'class': tf.io.FixedLenFeature([], tf.int64),\n    'id': tf.io.FixedLenFeature([], tf.string),\n    'image': tf.io.FixedLenFeature([], tf.string),\n}\n\ndef _parse_image_function(example_proto):\n  # Parse the input tf.Example proto using the dictionary above.\n  return tf.io.parse_single_example(example_proto, train_feature_description)\n\ntrain_ids = []\ntrain_class = []\ntrain_images = []\n\nfor i in train_files:\n  train_image_dataset = tf.data.TFRecordDataset(i)\n\n  train_image_dataset = train_image_dataset.map(_parse_image_function)\n\n  ids = [str(id_features['id'].numpy())[2:-1] for id_features in train_image_dataset] # [2:-1] is done to remove b' from 1st and 'from last in train id names\n  train_ids = train_ids + ids\n\n  classes = [int(class_features['class'].numpy()) for class_features in train_image_dataset]\n  train_class = train_class + classes\n\n  images = [image_features['image'].numpy() for image_features in train_image_dataset]\n  train_images = train_images + images","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Create a dictionary describing the features.\nval_feature_description = {\n    'class': tf.io.FixedLenFeature([], tf.int64),\n    'id': tf.io.FixedLenFeature([], tf.string),\n    'image': tf.io.FixedLenFeature([], tf.string),\n}\n\ndef _parse_image_function(example_proto):\n  # Parse the input tf.Example proto using the dictionary above.\n  return tf.io.parse_single_example(example_proto, val_feature_description)\n\nval_ids = []\nval_class = []\nval_images = []\n\nfor i in val_files:\n    val_image_dataset = tf.data.TFRecordDataset(i)\n\n    val_image_dataset = val_image_dataset.map(_parse_image_function)\n\n    ids = [str(image_features['id'].numpy())[2:-1] for image_features in val_image_dataset]\n    val_ids += ids\n\n    classes = [int(image_features['class'].numpy()) for image_features in val_image_dataset]\n    val_class += classes \n\n    images = [image_features['image'].numpy() for image_features in val_image_dataset]\n    val_images += images","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# dry run for testing\n\nimport IPython.display as display\n\ndisplay.display(display.Image(data=val_images[10000]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from PIL import Image\nimport cv2\nimport albumentations\nimport torch\nimport numpy as np\nimport io\nfrom torch.utils.data import Dataset\n\n# Making the dataset class for training and testing Flower images\n\nclass FlowerDataset(Dataset):\n    def __init__(self, id , classes , image , img_height , img_width, mean , std , is_valid):\n        self.id = id\n        self.classes = classes\n        self.image = image\n        self.is_valid = is_valid\n        if self.is_valid == 1: # transforms for validation images\n            self.aug = albumentations.Compose([\n               albumentations.Resize(img_height , img_width, always_apply = True) ,\n               albumentations.Normalize(mean , std , always_apply = True) \n            ])\n        else:                  # transfoms for training images \n            self.aug = albumentations.Compose([\n                albumentations.Resize(img_height , img_width, always_apply = True) ,\n                albumentations.Normalize(mean , std , always_apply = True),\n                albumentations.ShiftScaleRotate(shift_limit = 0.0625,\n                                                scale_limit = 0.1 ,\n                                                rotate_limit = 5,\n                                                p = 0.9)\n            ]) \n        \n    def __len__(self):\n        return len(self.id)\n    \n    def __getitem__(self, index):\n        id = self.id[index]\n        img = np.array(Image.open(io.BytesIO(self.image[index])))  # converting byte format of images to numpy array\n        img = cv2.resize(img, dsize=(224, 224), interpolation=cv2.INTER_CUBIC)\n        img = self.aug(image = img)['image']\n        img = np.transpose(img , (2,0,1)).astype(np.float32) # 2,0,1 because pytorch excepts image channel first then dimension of image\n       \n        if self.is_valid == 1: # converting the validation label dataset to categorical format so that it matches with lables of training data after applying Cutmix data augmentation \n            return torch.tensor(img, dtype = torch.float),torch.tensor(np.eye(104, dtype='float64')[int(self.classes[index])]) # 104 is the no. of classes\n        else:\n            return torch.tensor(img, dtype = torch.float),torch.tensor(self.classes[index], dtype = torch.long)\n        \n# creating object for the dataset class \ntrain_dataset = FlowerDataset(id = train_ids, classes = train_class, image = train_images, \n                        img_height = 224 , img_width = 224, \n                        mean = (0.485, 0.456, 0.406),\n                        std = (0.229, 0.224, 0.225) , is_valid = 0)\n\nval_dataset = FlowerDataset(id = val_ids, classes = val_class, image = val_images, \n                        img_height = 224 , img_width = 224, \n                        mean = (0.485, 0.456, 0.406),\n                        std = (0.229, 0.224, 0.225) , is_valid = 1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# sanity check for FlowerDataset class created\n\nimport matplotlib.pyplot as plt\n%matplotlib inline\n\nidx = 10000 # taking index for 10000th image out of 51000 images\nimg = val_dataset[idx][0]\n\nprint(val_dataset[idx][1]) # val_dataset label is one Hot encoded\n\nnpimg = img.numpy()\nplt.imshow(np.transpose(npimg, (1,2,0)))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# installing cutmix repository\n!pip install git+https://github.com/ildoonet/cutmix\n\nfrom cutmix.cutmix import CutMix\n\n# applying cut-mix to training images only\nCutMix_train_dataloader = CutMix(train_dataset, \n                          num_class=104, \n                          beta=1.0, \n                          prob=0.5, \n                          num_mix=2)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# creating a distributed sampler for training & validaion data\ntrain_sampler = torch.utils.data.distributed.DistributedSampler(\n          CutMix_train_dataloader,\n          num_replicas=xm.xrt_world_size(),\n          rank=xm.get_ordinal(),\n          shuffle=True)\n\nvalid_sampler = torch.utils.data.distributed.DistributedSampler(\n          val_dataset,\n          num_replicas=xm.xrt_world_size(),\n          rank=xm.get_ordinal(),\n          shuffle=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# setting up the validation data loader\n\nTRAIN_BATCH_SIZE = 128\n\nfrom torch.utils.data import DataLoader\n\ntraining_dataloader = DataLoader(CutMix_train_dataloader,\n                        num_workers=4,\n                        batch_size=TRAIN_BATCH_SIZE,\n                        sampler=train_sampler,\n                        drop_last=True\n                       )\n\nval_dataloader = DataLoader(val_dataset,\n                        num_workers=4,\n                        batch_size=TRAIN_BATCH_SIZE,\n                        sampler=valid_sampler,\n                        drop_last=False\n                       )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# creating object for integrating torch-xla with tpu \ndevice = xm.xla_device()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# downloading the pretrained model - efficientnet b7  \n\n!pip install efficientnet_pytorch\n\nimport efficientnet_pytorch\n\nmodel = efficientnet_pytorch.EfficientNet.from_pretrained('efficientnet-b0')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch.nn as nn\nimport torch.nn.functional as F\n\n# increasing few layers in our model\nclass EfficientNet_b0(nn.Module):\n    def __init__(self):\n        super(EfficientNet_b0, self).__init__()\n        self.model = efficientnet_pytorch.EfficientNet.from_pretrained('efficientnet-b0')\n        \n        self.classifier_layer = nn.Sequential(\n            nn.Linear(1280 , 512),\n            nn.BatchNorm1d(512),\n            nn.Dropout(0.2),\n            nn.Linear(512 , 256),\n            nn.Linear(256 , 104)\n        )\n        \n    def forward(self, inputs):\n        x = self.model.extract_features(inputs)\n\n        # Pooling and final linear layer\n        x = self.model._avg_pooling(x)\n        x = x.flatten(start_dim=1)\n        x = self.model._dropout(x)\n        x = self.classifier_layer(x)\n        return x\n    \nmodel = EfficientNet_b0()\nmodel = model.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# installing torchcontrib for Stochastic Weight Averaging in PyTorch \n!pip install torchcontrib","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#for Stochastic Weight Averaging in PyTorch\nfrom torchcontrib.optim import SWA\n\nfrom cutmix.utils import CutMixCrossEntropyLoss\n\nEPOCHS = 25\nnum_train_steps = int(len(train_dataset) / TRAIN_BATCH_SIZE / xm.xrt_world_size() * EPOCHS)\n\n# printing the no of training steps for each epoch of our training dataloader  \nxm.master_print(f'num_train_steps = {num_train_steps}, world_size={xm.xrt_world_size()}')\n\nbase_optimizer = torch.optim.Adam(model.classifier_layer.parameters(), lr=1e-4* xm.xrt_world_size())\n\noptimizer = SWA(base_optimizer, swa_start=5, swa_freq=5, swa_lr=0.05)\n\nloss_fn = CutMixCrossEntropyLoss(True)\n\nscheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# defining the training loop\ndef train_loop_fn(data_loader, model, optimizer, device, scheduler=None):\n    running_loss = 0.0\n    model.train()\n    for inputs,labels in data_loader:\n        inputs = inputs.to(device, dtype=torch.float)\n        labels = labels.to(device, dtype=torch.float)\n\n        optimizer.zero_grad()\n\n        outputs = model(inputs)\n        loss = loss_fn(outputs, labels)\n\n        loss.backward()\n        xm.optimizer_step(optimizer)\n\n        running_loss += loss.item() * inputs.size(0)\n\n    if scheduler is not None:\n        scheduler.step()\n            \n    train_loss = running_loss / float(len(train_dataset))\n    xm.master_print('training Loss: {:.4f}'.format(train_loss))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# defining the validation loop\ndef eval_loop_fn(data_loader, model, device):\n    running_loss = 0.0\n    running_corrects = 0.0\n    model.eval()\n    \n    for inputs,labels in data_loader:\n        inputs = inputs.to(device, dtype=torch.float)\n        labels = labels.to(device, dtype=torch.float)\n\n        outputs = model(inputs)\n        _, preds = torch.max(outputs, 1)\n        loss = loss_fn(outputs, labels)\n        \n        running_loss += loss.item() * inputs.size(0)\n        running_corrects += torch.sum(preds == torch.argmax(labels.data, dim=1))\n    \n    valid_loss = running_loss / float(len(val_dataset))\n    epoch_acc = running_corrects / float(len(val_dataset))\n    xm.master_print('validation Loss: {:.4f} Acc: {:.4f}'.format(valid_loss, epoch_acc))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# training the model in _run function\ndef _run():\n    for param in model.parameters():\n        param.requires_grad = False\n    \n    for param in model.classifier_layer.parameters():\n        param.requires_grad = True\n    \n    for epoch in range(EPOCHS):\n        xm.master_print(f\"Epoch --> {epoch+1} / {EPOCHS}\")\n        xm.master_print(f\"-------------------------------\")\n        para_loader = pl.ParallelLoader(training_dataloader, [device])\n        train_loop_fn(para_loader.per_device_loader(device), model, optimizer, device, scheduler=scheduler)\n\n        para_loader = pl.ParallelLoader(val_dataloader, [device])\n        eval_loop_fn(para_loader.per_device_loader(device), model, device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# initializing the training of model\ndef _mp_fn(rank, flags):\n    torch.set_default_tensor_type('torch.FloatTensor')\n    a = _run()\n    optimizer.swap_swa_sgd()\n    \n# applying multiprocessing so that images get paralley trained in different cores of kaggle-tpu\nFLAGS={}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=1, 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}