{"cells":[{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"%reload_ext autoreload\n%autoreload 2\n%matplotlib inline","execution_count":null,"outputs":[]},{"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":{"trusted":true},"cell_type":"code","source":"import glob\nimport io\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport os\nimport pandas as pd\nimport random\nfrom PIL import Image\nimport time\n#from tqdm.notebook import tqdm\n\nimport tensorflow as tf\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, datasets\n\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.utils.utils as xu\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Read tf-record for Pytorch\n\nBased on https://medium.com/analytics-vidhya/how-to-read-tfrecords-files-in-pytorch-72763786743f"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_files = glob.glob(\"../input/tpu-getting-started/tfrecords-jpeg-224x224/train/*.tfrec\")\nval_files = glob.glob(\"../input/tpu-getting-started/tfrecords-jpeg-224x224/val/*.tfrec\")\ntest_files = glob.glob(\"../input/tpu-getting-started/tfrecords-jpeg-224x224/test/*.tfrec\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_feature_description = {\n    'class': tf.io.FixedLenFeature([], tf.int64),\n    'id': tf.io.FixedLenFeature([], tf.string),\n    'image': tf.io.FixedLenFeature([], tf.string),\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _parse_image_function(example_proto):\n    return tf.io.parse_single_example(example_proto, train_feature_description)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_ids = []\ntrain_class = []\ntrain_images = []\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]\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":"val_ids = []\nval_class = []\nval_images = []\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(id_features['id'].numpy())[2:-1] for id_features in val_image_dataset]\n    val_ids = val_ids + ids\n\n    classes = [int(class_features['class'].numpy()) for class_features in val_image_dataset]\n    val_class = val_class + classes\n\n    images = [image_features['image'].numpy() for image_features in val_image_dataset]\n    val_images = val_images + images","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_feature_description = {\n    'id': tf.io.FixedLenFeature([], tf.string),\n    'image': tf.io.FixedLenFeature([], tf.string),\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _parse_image_function_test(example_proto):\n    return tf.io.parse_single_example(example_proto, test_feature_description)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_ids = []\ntest_images = []\nfor i in test_files:\n    test_image_dataset = tf.data.TFRecordDataset(i)\n    \n    test_image_dataset = test_image_dataset.map(_parse_image_function_test)\n\n    ids = [str(id_features['id'].numpy())[2:-1] for id_features in test_image_dataset]\n    test_ids = test_ids + ids\n\n    images = [image_features['image'].numpy() for image_features in test_image_dataset]\n    test_images = test_images + images","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"import cv2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class FlowerDS():\n    def __init__(self, ids, cls, imgs, transforms, is_test=False):\n        self.ids = ids\n        if not is_test:\n            self.cls = cls\n        self.imgs = imgs\n        self.transforms = transforms\n        self.is_test = is_test\n    \n    def __len__(self):\n        return len(self.ids)\n\n    def __getitem__(self, idx):\n        img = self.imgs[idx]\n        img = Image.open(io.BytesIO(img))\n        img = self.transforms(img)\n        if self.is_test:\n            return img, -1, self.ids[idx]\n        return img, int(self.cls[idx]), self.ids[idx]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# normalize stats"},{"metadata":{"trusted":true},"cell_type":"code","source":"mean = [0.485, 0.456, 0.406]\nstd = [0.229, 0.224, 0.225]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"normalize = transforms.Normalize(mean=mean, std=std)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Vision Transformer"},{"metadata":{},"cell_type":"markdown","source":"We'll try vision transformer from recent paper https://arxiv.org/abs/2010.11929 <br>\nThe implementation can be found @ https://github.com/nachiket273/VisTrans<br>\nI have created simple library for the same and it can be installed using<br>\npip install vistrans<br>\nFurther info can be found @ https://pypi.org/project/vistrans/"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip uninstall vistrans -y\n!pip install vistrans","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from vistrans import VisionTransformer","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def save_checkpoint(model, is_best, filename='./checkpoint.pth'):\n    \"\"\"Save checkpoint if a new best is achieved\"\"\"\n    if is_best:\n        xm.save(model.state_dict(), filename)  # save checkpoint\n    else:\n        print (\"=> Validation Accuracy did not improve\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_checkpoint(model, filename = './checkpoint.pth'):\n    sd = torch.load(filename, map_location=lambda storage, loc: storage)\n    names = set(model.state_dict().keys())\n    for n in list(sd.keys()):\n        if n not in names and n+'_raw' in names:\n            if n+'_raw' not in sd: sd[n+'_raw'] = sd[n]\n            del sd[n]\n    model.load_state_dict(sd)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model and load weights"},{"metadata":{"trusted":true},"cell_type":"code","source":"VisionTransformer.list_pretrained()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_model(name ='vit_b16_224'):\n    model = VisionTransformer.create_pretrained(name, num_classes=104)\n    for param in model.parameters():\n        param.require_grad = True\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"SERIAL_EXEC = xmp.MpSerialExecutor()\nWRAPPED_MODEL = xmp.MpModelWrapper(get_model())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Stats"},{"metadata":{"trusted":true},"cell_type":"code","source":"class AvgStats(object):\n    def __init__(self):\n        self.reset()\n        \n    def reset(self):\n        self.losses =[]\n        self.precs =[]\n        self.its = []\n        \n    def append(self, loss, prec, it):\n        self.losses.append(loss)\n        self.precs.append(prec)\n        self.its.append(it)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trn_stat = AvgStats()\nval_stat = AvgStats()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Fit"},{"metadata":{"trusted":true},"cell_type":"code","source":"def fit(rank, flags):\n    torch.set_default_tensor_type('torch.FloatTensor')\n    bs = flags['bs']\n    epochs = flags['epochs']\n    WRAPPED_MODEL = flags['model']\n    torch.manual_seed(719)\n    device = xm.xla_device()\n    \n    def get_dataset():\n        train_transforms = transforms.Compose([\n                        transforms.RandomResizedCrop(224),\n                        transforms.RandomHorizontalFlip(),\n                        transforms.RandomVerticalFlip(),\n                        transforms.ToTensor(),\n                        normalize,\n                        transforms.RandomErasing()\n                    ])\n\n        test_transforms = transforms.Compose([\n                        transforms.CenterCrop(224),\n                        transforms.Resize(224),\n                        transforms.ToTensor(),\n                        normalize\n                    ])\n\n        train_ds = FlowerDS(train_ids, train_class, train_images, train_transforms)\n        valid_ds = FlowerDS(val_ids, val_class, val_images, test_transforms)\n\n        return train_ds, valid_ds\n    \n    train_ds, valid_ds = SERIAL_EXEC.run(get_dataset)\n    \n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n        train_ds,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True\n    )\n\n    valid_sampler = torch.utils.data.distributed.DistributedSampler(\n        valid_ds,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True\n    )\n    \n    train_loader = DataLoader(train_ds, bs, sampler=train_sampler, num_workers=1, pin_memory=True)\n    valid_loader = DataLoader(valid_ds, bs, sampler=valid_sampler, num_workers=1, pin_memory=True)\n    \n    \n    model = WRAPPED_MODEL.to(device)\n    \n    criterion = nn.CrossEntropyLoss()\n    optimizer = torch.optim.SGD(model.parameters(), lr=3e-2*xm.xrt_world_size(), momentum=0.9)\n    \n    def train(loader, epoch, model, optimizer, criterion):\n        #tracker = xm.RateTracker()\n        model.train()\n        running_loss = 0.\n        running_acc = 0.\n        tot = 0\n        start_time = time.time()\n        for i, (ip, tgt, _) in enumerate(loader):\n            #ip, tgt = ip.to(device), tgt.to(device)                            \n            output = model(ip)\n            loss = criterion(output, tgt)\n            running_loss += loss.item()\n            tot += ip.shape[0]\n\n            # Append outputs\n            _, pred = output.max(dim=1)\n            running_acc += torch.sum(pred == tgt.data)\n\n            # compute gradient and do SGD step\n            optimizer.zero_grad()\n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            #optimizer.step()\n            xm.optimizer_step(optimizer)\n\n        trn_time = time.time() - start_time        \n        trn_acc = (running_acc/tot) * 100\n        trn_loss = running_loss/len(loader)\n        return trn_acc, trn_loss, trn_time\n    \n    def test(model, loader, criterion):\n        with torch.no_grad():\n            model.eval()\n            running_loss = 0.\n            running_acc = 0.\n            tot = 0\n            start_time = time.time()\n            for i, (ip, tgt, _) in enumerate(loader):\n                #ip, tgt = ip.to(device), tgt.to(device)\n                output = model(ip)\n                loss = criterion(output, tgt)\n                running_loss += loss.item()\n                tot += ip.shape[0]\n                _, pred = output.max(dim=1)\n                running_acc += torch.sum(pred == tgt.data)\n\n            val_time = time.time() - start_time\n            val_acc = (running_acc/tot) * 100\n            val_loss = running_loss/len(loader)\n            return val_acc, val_loss, val_time\n        \n    sched = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs, 3e-5)\n    #sched = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[7, 14, 20])\n    for j in range(1, epochs+1):\n        para_loader = pl.ParallelLoader(train_loader, [device])\n        trn_acc, trn_losses, trn_time = train(para_loader.per_device_loader(device), j, model,\n                                             optimizer, criterion)\n        trn_stat.append(trn_losses, trn_acc, trn_time)\n        para_loader = pl.ParallelLoader(valid_loader, [device])\n        val_acc, val_losses, val_time = test(model, para_loader.per_device_loader(device), criterion)\n        val_stat.append(val_losses, val_acc, val_time)            \n        sched.step()\n        print(\"Epoch::{}, Trn_loss::{:06.8f}, Val_loss::{:06.8f}, Trn_F1::{:06.8f}, Val_F1::{:06.8f}\"\n            .format(j, trn_losses, val_losses, trn_acc, val_acc))\n        \n    save_checkpoint(model, True, './best_model.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"flags = dict()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"flags['epochs'] = 25\nflags['bs'] = 32\nflags['model'] = WRAPPED_MODEL","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":false},"cell_type":"code","source":"xmp.spawn(fit, args=(flags,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!mv best_model.pth best_model_vit_b16_224.pth","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"WRAPPED_MODEL1 = xmp.MpModelWrapper(get_model('vit_l16_224'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"flags['bs'] = 16\nflags['model'] = WRAPPED_MODEL1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"xmp.spawn(fit, args=(flags,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_transforms = transforms.Compose([\n                        transforms.CenterCrop(224),\n                        transforms.Resize(224),\n                        transforms.ToTensor(),\n                        normalize\n                    ])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_ds = FlowerDS(test_ids, [], test_images, test_transforms, True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = xm.xla_device()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"testloader = DataLoader(test_ds, 16, num_workers=4, pin_memory=True, shuffle=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def predict(loader, device):\n    with torch.no_grad():\n        torch.cuda.empty_cache()\n        model.eval()\n        model1.eval()\n        preds = dict()\n        for i, (ip, _, ids) in enumerate(loader):\n            ip = ip.to(device)\n            output = model(ip)\n            _, pred = output.max(dim=1)\n            output1 = model1(ip)\n            _, pred1 = output1.max(dim=1)\n            for i, j, k in zip(ids, pred.cpu().detach(), pred1.cpu().detach()):\n                preds[i] = int((j.item() + k.item())/2)\n        return preds","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = WRAPPED_MODEL1.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"load_checkpoint(model, './best_model.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model1 = WRAPPED_MODEL.to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"load_checkpoint(model1, './best_model_vit_b16_224.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"preds = predict(testloader, device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import csv","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sub_csv = pd.read_csv('../input/tpu-getting-started/sample_submission.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sub_csv.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for key in preds.keys():\n    sub_csv.loc[sub_csv['id'] == key, 'label'] = preds[key]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sub_csv.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sub_csv.to_csv('submission.csv', index=False)","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}