{"cells":[{"metadata":{},"cell_type":"markdown","source":"

Vision Transformers are Easy

\n
\n
"},{"metadata":{},"cell_type":"markdown","source":"

Please give some due credit to Abhinand for helping make this Kernel a reality.

"},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.execute_input":"2020-12-02T18:32:11.628947Z","iopub.status.busy":"2020-12-02T18:32:11.628076Z","iopub.status.idle":"2020-12-02T18:32:24.574113Z","shell.execute_reply":"2020-12-02T18:32:24.573413Z"},"papermill":{"duration":12.981293,"end_time":"2020-12-02T18:32:24.574247","exception":false,"start_time":"2020-12-02T18:32:11.592954","status":"completed"},"tags":[],"trusted":true,"collapsed":true},"cell_type":"code","source":"!cp -r ../input/vittutorialillustrations/* ./ \n\n!pip install nb_black\n%load_ext nb_black","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

Installing Pytorch-XLA inorder to access the TPU:

"},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.execute_input":"2020-12-02T18:32:25.025299Z","iopub.status.busy":"2020-12-02T18:32:25.024308Z","iopub.status.idle":"2020-12-02T18:34:10.700354Z","shell.execute_reply":"2020-12-02T18:34:10.69964Z"},"papermill":{"duration":105.718607,"end_time":"2020-12-02T18:34:10.7005","exception":false,"start_time":"2020-12-02T18:32:24.981893","status":"completed"},"tags":[],"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 1.7\n!pip install timm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

Importing some really important stuff in here. I hope you are still awake as you course through this notebook.

"},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:10.826084Z","iopub.status.busy":"2020-12-02T18:34:10.824931Z","iopub.status.idle":"2020-12-02T18:34:13.015202Z","shell.execute_reply":"2020-12-02T18:34:13.015867Z"},"papermill":{"duration":2.258623,"end_time":"2020-12-02T18:34:13.016037","exception":false,"start_time":"2020-12-02T18:34:10.757414","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"# Numpy, Pandas and MatPlot\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\n\nplt.style.use(\"ggplot\")\n\n# Who doesn't love PyTorch\nimport torch\nimport torch.nn as nn\n\n# For the Transformations\nimport torchvision.transforms as transforms\n\n# For the TPU to work\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.distributed.parallel_loader as pl\n\nimport timm\n\n# Some Native Python Libraries\nimport gc\nimport os\nimport time\nimport random\nfrom datetime import datetime\n\nfrom PIL import Image\nfrom tqdm.notebook import tqdm\nfrom sklearn import model_selection, metrics","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

🚨 Jargon Alert 🚨

I know, TPU Parallelization sounds aweful but you can watch this video about it. Watch till 1:04. It is by folks at TensorFlow so you need not watch it for a very long time.
(No one likes TensorFlow or atleast in the world I live in)

"},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:13.138754Z","iopub.status.busy":"2020-12-02T18:34:13.137657Z","iopub.status.idle":"2020-12-02T18:34:13.146711Z","shell.execute_reply":"2020-12-02T18:34:13.145929Z"},"papermill":{"duration":0.073227,"end_time":"2020-12-02T18:34:13.146842","exception":false,"start_time":"2020-12-02T18:34:13.073615","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"# For parallelization in TPUs\nos.environ[\"XLA_USE_BF16\"] = \"1\"\nos.environ[\"XLA_TENSOR_ALLOCATOR_MAXSIZE\"] = \"100000000\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

When we say that:
\"Hey let's seed this!\"
\nDL Developers aren't really thinking about gardening, or well, uh. It has to do maybe a little bit about the reproducebility of the model.

Check out this Medium Article:
Properly Setting the Random Seed in ML Experiments. Not as Simple as You Might Imagine 🐵
to get a better idea.

"},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:13.273861Z","iopub.status.busy":"2020-12-02T18:34:13.273082Z","iopub.status.idle":"2020-12-02T18:34:13.288963Z","shell.execute_reply":"2020-12-02T18:34:13.288196Z"},"papermill":{"duration":0.083797,"end_time":"2020-12-02T18:34:13.289095","exception":false,"start_time":"2020-12-02T18:34:13.205298","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"def seed_everything(seed):\n \"\"\"\n Seeds basic parameters for reproductibility of results\n \n Arguments:\n seed {int} -- Number of the seed\n \"\"\"\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 torch.backends.cudnn.benchmark = False\n\n\nseed_everything(1001)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

Specifying necessary paths and other necessities:

"},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:13.41906Z","iopub.status.busy":"2020-12-02T18:34:13.418234Z","iopub.status.idle":"2020-12-02T18:34:13.429184Z","shell.execute_reply":"2020-12-02T18:34:13.428386Z"},"papermill":{"duration":0.080249,"end_time":"2020-12-02T18:34:13.429314","exception":false,"start_time":"2020-12-02T18:34:13.349065","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"# general global variables\nDATA_PATH = \"../input/cassava-leaf-disease-classification\"\nTRAIN_PATH = \"../input/cassava-leaf-disease-classification/train_images/\"\nTEST_PATH = \"../input/cassava-leaf-disease-classification/test_images/\"\nMODEL_PATH = (\n \"../input/vit-base-models-pretrained-pytorch/jx_vit_base_p16_224-80ecf9dd.pth\"\n)\n\n# model specific global variables\n\n# THIS IS THE IMAGE SIZE IN PIXELS:\nIMG_SIZE = 224\n\nBATCH_SIZE = 16\n\nLR = 2e-05\n\nN_EPOCHS = 10","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:13.557875Z","iopub.status.busy":"2020-12-02T18:34:13.557096Z","iopub.status.idle":"2020-12-02T18:34:13.603956Z","shell.execute_reply":"2020-12-02T18:34:13.603128Z"},"papermill":{"duration":0.114526,"end_time":"2020-12-02T18:34:13.604087","exception":false,"start_time":"2020-12-02T18:34:13.489561","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"# Lets take a look at the data:\ndf = pd.read_csv(os.path.join(DATA_PATH, \"train.csv\"))\n\n# Splitting the dataset into training and testsets:\ntrain_df, valid_df = model_selection.train_test_split(\n df, test_size=0.1, random_state=42, stratify=df.label.values\n)\n\ndf.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

We will now be visualizing on how\n

Plotting to see how various labels are distributed:

"},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:13.889294Z","iopub.status.busy":"2020-12-02T18:34:13.888176Z","iopub.status.idle":"2020-12-02T18:34:14.103989Z","shell.execute_reply":"2020-12-02T18:34:14.103314Z"},"papermill":{"duration":0.288382,"end_time":"2020-12-02T18:34:14.104118","exception":false,"start_time":"2020-12-02T18:34:13.815736","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"df.label.value_counts().plot(kind=\"bar\")","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:14.240631Z","iopub.status.busy":"2020-12-02T18:34:14.239815Z","iopub.status.idle":"2020-12-02T18:34:14.272019Z","shell.execute_reply":"2020-12-02T18:34:14.271271Z"},"papermill":{"duration":0.102979,"end_time":"2020-12-02T18:34:14.272159","exception":false,"start_time":"2020-12-02T18:34:14.16918","status":"completed"},"tags":[]},"cell_type":"markdown","source":"

Defining the CasavaDataset:

"},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:14.41932Z","iopub.status.busy":"2020-12-02T18:34:14.418142Z","iopub.status.idle":"2020-12-02T18:34:14.444009Z","shell.execute_reply":"2020-12-02T18:34:14.443354Z"},"papermill":{"duration":0.106308,"end_time":"2020-12-02T18:34:14.444148","exception":false,"start_time":"2020-12-02T18:34:14.33784","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"class CassavaDataset(torch.utils.data.Dataset):\n \"\"\"\n Helper Class to create the pytorch dataset\n \"\"\"\n\n def __init__(self, df, data_path=DATA_PATH, mode=\"train\", transforms=None):\n super().__init__()\n self.df_data = df.values\n self.data_path = data_path\n self.transforms = transforms\n self.mode = mode\n self.data_dir = \"train_images\" if mode == \"train\" else \"test_images\"\n\n def __len__(self):\n return len(self.df_data)\n\n def __getitem__(self, index):\n img_name, label = self.df_data[index]\n img_path = os.path.join(self.data_path, self.data_dir, img_name)\n img = Image.open(img_path).convert(\"RGB\")\n\n if self.transforms is not None:\n image = self.transforms(img)\n\n return image, label","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

Data Augmentation to go through:

\n

We go through some really particular Data Augmentations. I have to say that our dude Abhinand did an amazing job here. Kudos to him.

\nIf you want a better explanation head to the docs of torchvision.transforms here.

"},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:14.590922Z","iopub.status.busy":"2020-12-02T18:34:14.589829Z","iopub.status.idle":"2020-12-02T18:34:14.612461Z","shell.execute_reply":"2020-12-02T18:34:14.61172Z"},"papermill":{"duration":0.101631,"end_time":"2020-12-02T18:34:14.612584","exception":false,"start_time":"2020-12-02T18:34:14.510953","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"# Time to do some data augmentations.\n\ntransforms_train = transforms.Compose(\n [\n # Resizing the image to the previously stated size\n transforms.Resize((IMG_SIZE, IMG_SIZE)),\n # Randomly flipping the images horizontally with the probability of 30%\n transforms.RandomHorizontalFlip(p=0.3),\n # Randomly flipping the images vertically with the probability of 30%\n transforms.RandomVerticalFlip(p=0.3),\n # Randomly Rotating the images by 10 degrees\n transforms.RandomRotation(10),\n # Randomly 10 degrees worth of Affine\n transforms.RandomAffine(10),\n # Croppping the images to the stated\n transforms.RandomResizedCrop(IMG_SIZE),\n # Converting the image to tensor\n transforms.ToTensor(),\n # Normalizing\n transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),\n ]\n)\n\ntransforms_valid = transforms.Compose(\n [\n transforms.Resize((IMG_SIZE, IMG_SIZE)),\n transforms.ToTensor(),\n transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),\n ]\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

Defining the Model

"},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:14.936128Z","iopub.status.busy":"2020-12-02T18:34:14.924554Z","iopub.status.idle":"2020-12-02T18:34:15.031764Z","shell.execute_reply":"2020-12-02T18:34:15.030784Z"},"papermill":{"duration":0.197739,"end_time":"2020-12-02T18:34:15.031965","exception":false,"start_time":"2020-12-02T18:34:14.834226","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"class ViTBase16(nn.Module):\n def __init__(self, n_classes, pretrained=False):\n\n super(ViTBase16, self).__init__()\n\n self.model = timm.create_model(\"vit_base_patch16_224\", pretrained=False)\n\n # I mean it is pretrained?!\n if pretrained:\n self.model.load_state_dict(torch.load(MODEL_PATH))\n\n self.model.head = nn.Linear(self.model.head.in_features, n_classes)\n\n def forward(self, x):\n x = self.model(x)\n return x\n\n def train_one_epoch(self, train_loader, criterion, optimizer, device):\n\n # The Accuracy starts from Zero:\n epoch_loss = 0.0\n epoch_accuracy = 0.0\n\n # This bad boy trains the model\n self.model.train()\n for i, (data, target) in enumerate(train_loader):\n\n # Clear the gradients of all optimized variables\n optimizer.zero_grad()\n # forward pass: compute predicted outputs by passing inputs to the model\n output = self.forward(data)\n # calculate the batch loss\n loss = criterion(output, target)\n # backward pass: compute gradient of the loss with respect to model parameters\n loss.backward()\n # Calculate Accuracy\n accuracy = (output.argmax(dim=1) == target).float().mean()\n # update training loss and accuracy\n epoch_loss += loss\n epoch_accuracy += accuracy\n\n # perform a single optimization step (parameter update)\n if device.type == \"xla\":\n xm.optimizer_step(optimizer)\n\n if i % 20 == 0:\n xm.master_print(f\"\\tBATCH {i+1}/{len(train_loader)} - LOSS: {loss}\")\n\n else:\n optimizer.step()\n\n return epoch_loss / len(train_loader), epoch_accuracy / len(train_loader)\n\n def validate_one_epoch(self, valid_loader, criterion, device):\n # keep track of validation loss\n valid_loss = 0.0\n valid_accuracy = 0.0\n\n ######################\n # validate the model #\n ######################\n self.model.eval()\n for data, target in valid_loader:\n # move tensors to GPU if CUDA is available\n if device.type == \"cuda\":\n data, target = data.cuda(), target.cuda()\n elif device.type == \"xla\":\n data = data.to(device, dtype=torch.float32)\n target = target.to(device, dtype=torch.int64)\n\n with torch.no_grad():\n # forward pass: compute predicted outputs by passing inputs to the model\n output = self.model(data)\n # calculate the batch loss\n loss = criterion(output, target)\n # Calculate Accuracy\n accuracy = (output.argmax(dim=1) == target).float().mean()\n # update average validation loss and accuracy\n valid_loss += loss\n valid_accuracy += accuracy\n\n return valid_loss / len(valid_loader), valid_accuracy / len(valid_loader)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

Fitting the Model to the TPU 🏗

"},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:15.199425Z","iopub.status.busy":"2020-12-02T18:34:15.198635Z","iopub.status.idle":"2020-12-02T18:34:15.23679Z","shell.execute_reply":"2020-12-02T18:34:15.236128Z"},"papermill":{"duration":0.131593,"end_time":"2020-12-02T18:34:15.236935","exception":false,"start_time":"2020-12-02T18:34:15.105342","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"def fit_tpu(\n model, epochs, device, criterion, optimizer, train_loader, valid_loader=None\n):\n\n valid_loss_min = np.Inf # track change in validation loss\n\n # keeping track of losses as it happen\n train_losses = []\n valid_losses = []\n train_accs = []\n valid_accs = []\n\n for epoch in range(1, epochs + 1):\n gc.collect()\n para_train_loader = pl.ParallelLoader(train_loader, [device])\n\n xm.master_print(f\"{'='*50}\")\n xm.master_print(f\"EPOCH {epoch} - TRAINING...\")\n train_loss, train_acc = model.train_one_epoch(\n para_train_loader.per_device_loader(device), criterion, optimizer, device\n )\n xm.master_print(\n f\"\\n\\t[TRAIN] EPOCH {epoch} - LOSS: {train_loss}, ACCURACY: {train_acc}\\n\"\n )\n train_losses.append(train_loss)\n train_accs.append(train_acc)\n gc.collect()\n\n if valid_loader is not None:\n gc.collect()\n para_valid_loader = pl.ParallelLoader(valid_loader, [device])\n xm.master_print(f\"EPOCH {epoch} - VALIDATING...\")\n valid_loss, valid_acc = model.validate_one_epoch(\n para_valid_loader.per_device_loader(device), criterion, device\n )\n xm.master_print(f\"\\t[VALID] LOSS: {valid_loss}, ACCURACY: {valid_acc}\\n\")\n valid_losses.append(valid_loss)\n valid_accs.append(valid_acc)\n gc.collect()\n\n # save model if validation loss has decreased\n if valid_loss <= valid_loss_min and epoch != 1:\n xm.master_print(\n \"Validation loss decreased ({:.4f} --> {:.4f}). Saving model ...\".format(\n valid_loss_min, valid_loss\n )\n )\n # xm.save(model.state_dict(), 'best_model.pth')\n\n valid_loss_min = valid_loss\n\n return {\n \"train_loss\": train_losses,\n \"valid_losses\": valid_losses,\n \"train_acc\": train_accs,\n \"valid_acc\": valid_accs,\n }","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

Grabbing the Model

"},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:15.386773Z","iopub.status.busy":"2020-12-02T18:34:15.385757Z","iopub.status.idle":"2020-12-02T18:34:17.883884Z","shell.execute_reply":"2020-12-02T18:34:17.88284Z"},"papermill":{"duration":2.57545,"end_time":"2020-12-02T18:34:17.884083","exception":false,"start_time":"2020-12-02T18:34:15.308633","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"model = ViTBase16(n_classes=5, pretrained=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

Training the Model

"},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:18.110856Z","iopub.status.busy":"2020-12-02T18:34:18.107691Z","iopub.status.idle":"2020-12-02T18:34:18.154359Z","shell.execute_reply":"2020-12-02T18:34:18.154988Z"},"papermill":{"duration":0.156612,"end_time":"2020-12-02T18:34:18.15515","exception":false,"start_time":"2020-12-02T18:34:17.998538","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"def _run():\n train_dataset = CassavaDataset(train_df, transforms=transforms_train)\n valid_dataset = CassavaDataset(valid_df, transforms=transforms_valid)\n\n train_sampler = torch.utils.data.distributed.DistributedSampler(\n train_dataset,\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_dataset,\n num_replicas=xm.xrt_world_size(),\n rank=xm.get_ordinal(),\n shuffle=False,\n )\n\n train_loader = torch.utils.data.DataLoader(\n dataset=train_dataset,\n batch_size=BATCH_SIZE,\n sampler=train_sampler,\n drop_last=True,\n num_workers=8,\n )\n\n valid_loader = torch.utils.data.DataLoader(\n dataset=valid_dataset,\n batch_size=BATCH_SIZE,\n sampler=valid_sampler,\n drop_last=True,\n num_workers=8,\n )\n\n criterion = nn.CrossEntropyLoss()\n # device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n device = xm.xla_device()\n model.to(device)\n\n lr = LR * xm.xrt_world_size()\n optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n\n xm.master_print(f\"INITIALIZING TRAINING ON {xm.xrt_world_size()} TPU CORES\")\n start_time = datetime.now()\n xm.master_print(f\"Start Time: {start_time}\")\n\n logs = fit_tpu(\n model=model,\n epochs=N_EPOCHS,\n device=device,\n criterion=criterion,\n optimizer=optimizer,\n train_loader=train_loader,\n valid_loader=valid_loader,\n )\n\n xm.master_print(f\"Execution time: {datetime.now() - start_time}\")\n\n xm.master_print(\"Saving Model\")\n xm.save(\n model.state_dict(), f'model_5e_{datetime.now().strftime(\"%Y%m%d-%H%M\")}.pth'\n )","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-12-02T18:34:18.313245Z","iopub.status.busy":"2020-12-02T18:34:18.312128Z","iopub.status.idle":"2020-12-02T19:08:15.586913Z","shell.execute_reply":"2020-12-02T19:08:15.587804Z"},"papermill":{"duration":2037.357951,"end_time":"2020-12-02T19:08:15.588226","exception":false,"start_time":"2020-12-02T18:34:18.230275","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"# Start training processes\ndef _mp_fn(rank, flags):\n torch.set_default_tensor_type(\"torch.FloatTensor\")\n a = _run()\n\n\n# _run()\nFLAGS = {}\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method=\"fork\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"

Well folks, that would be it, I'll update this every once in a while to see if there's more scope for improvement. I will not be using this notebook to make submissions because parts of it are code that were adapted from other's works. Please give some due credit to Abhinand too.

Thanks for going through the notebook. An upvote would be appreciated 😋

"},{"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}