{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Cassava PyTorch XLA/TPU starter (GPU inference)\n\n### If you found this helpful, please give it an upvote!\n\nThis is a GPU inference kernel for my [PyTorch TPU starter kernel](https://www.kaggle.com/tanlikesmath/cassava-pytorch-xla-tpu-starter-training) (TPU notebook inference is not allowed in this competition)."},{"metadata":{},"cell_type":"markdown","source":"## Installs & Imports"},{"metadata":{},"cell_type":"markdown","source":"The below cell will install the [timm]() library, which is what we will use to define our models and get pretrained weights."},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install ../input/pytorch-image-models/timm-0.3.1-py3-none-any.whl > /dev/null","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here are all of our imports!"},{"metadata":{"trusted":true},"cell_type":"code","source":"import gc\nimport os\nimport time\nimport torch\nimport albumentations\n\nimport numpy as np\nimport pandas as pd\n\nimport cv2\nfrom PIL import Image\n\nimport torch.nn as nn\nfrom sklearn import metrics\nfrom sklearn import model_selection\nfrom torch.nn import functional as F\nfrom torch.optim import Adam\n\nimport timm\n\nfrom tqdm.notebook import tqdm\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Definitions\n\nNow let's define the necessary functions and variables needed for training."},{"metadata":{},"cell_type":"markdown","source":"These are the flags for inference."},{"metadata":{"trusted":true},"cell_type":"code","source":"FLAGS = {\n    \n    'num_folds': 5,\n    'model': 'resnext50_32x4d',\n    'model_path': '../input/cassava-pytorch-xla-tpu-starter-training/',\n    'batch_size': 32,\n    'epochs': 10,\n    'num_workers': 4,\n}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here, I define a model class for the timm models."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Using Ross Wightman's timm package\nclass TimmModels(nn.Module):\n    def __init__(self, model_name,pretrained=True, num_classes=5):\n        super(TimmModels, self).__init__()\n        self.m = timm.create_model(model_name,pretrained=pretrained)\n        model_list = list(self.m.children())\n        model_list[-1] = nn.Linear(\n            in_features=model_list[-1].in_features, \n            out_features=num_classes, \n            bias=True\n        )\n        self.m = nn.Sequential(*model_list)\n        \n    def forward(self, image):\n        out = self.m(image)\n        return out\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here, I define a class for the PyTorch Dataset (taken from @abhishek's amazing [Tez package](https://github.com/abhishekkrthakur/tez))."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Image Dataset class taken from Abhishek's tez package\n\nclass ImageDataset:\n    def __init__(\n        self,\n        image_paths,\n        targets,\n        resize,\n        augmentations=None,\n        backend=\"pil\",\n        channel_first=True,\n    ):\n        \"\"\"\n        :param image_paths: list of paths to images\n        :param targets: numpy array\n        :param resize: tuple or None\n        :param augmentations: albumentations augmentations\n        \"\"\"\n        self.image_paths = image_paths\n        self.targets = targets\n        self.resize = resize\n        self.augmentations = augmentations\n        self.backend = backend\n        self.channel_first = channel_first\n\n    def __len__(self):\n        return len(self.image_paths)\n\n    def __getitem__(self, item):\n        targets = self.targets[item]\n        if self.backend == \"pil\":\n            image = Image.open(self.image_paths[item])\n            if self.resize is not None:\n                image = image.resize(\n                    (self.resize[1], self.resize[0]), resample=Image.BILINEAR\n                )\n            image = np.array(image)\n            if self.augmentations is not None:\n                augmented = self.augmentations(image=image)\n                image = augmented[\"image\"]\n        elif self.backend == \"cv2\":\n            image = cv2.imread(self.image_paths[item])\n            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n            if self.resize is not None:\n                image = cv2.resize(\n                    image,\n                    (self.resize[1], self.resize[0]),\n                    interpolation=cv2.INTER_CUBIC,\n                )\n            if self.augmentations is not None:\n                augmented = self.augmentations(image=image)\n            image = augmented[\"image\"]\n        else:\n            raise Exception(\"Backend not implemented\")\n        if self.channel_first:\n            image = np.transpose(image, (2, 0, 1)).astype(np.float32)\n        return {\n            \"image\": torch.tensor(image),\n            \"targets\": torch.tensor(targets),\n        }","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Inference code"},{"metadata":{},"cell_type":"markdown","source":"Let's start predicting! To do so, we start by initializing the model."},{"metadata":{"trusted":true},"cell_type":"code","source":"MX = TimmModels(FLAGS['model'],pretrained=False, num_classes=5)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's now define our inference function."},{"metadata":{"trusted":true},"cell_type":"code","source":"def single_model_inference_fn(data_loader, model, device):\n    fin_outputs = []\n    model.eval()\n    model.to(device)\n    for bi, d in enumerate(tqdm(data_loader)): # enumerate through dataloader\n        \n        images = d['image'].to(device) # obtain the images\n\n        # pass image to model\n        outputs = model(images)\n\n        # Add the outputs and targets to a list \n        outputs_np = outputs.cpu().detach().numpy().tolist()\n        fin_outputs.extend(outputs_np)    \n        del outputs_np\n        gc.collect() # delete for memory conservation\n                \n    o = np.array(fin_outputs)\n    return o","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def tta(num_times, data_loader, model, device):\n    final_preds = None\n    for i in range(num_times):\n        temp_preds = single_model_inference_fn(data_loader, model, device)\n        if final_preds is None:\n            final_preds = temp_preds\n        else:\n            final_preds += temp_preds\n        \n    final_preds /= num_times\n    return final_preds","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda:0')\n\n\nmean = (0.485, 0.456, 0.406)\nstd = (0.229, 0.224, 0.225)\n\ntest_aug = albumentations.Compose(\n    [\n        albumentations.Resize(256, 256, p=1.0),\n        albumentations.Normalize(\n            mean, \n            std, \n            max_pixel_value=255.0, \n            always_apply=True\n        ),\n        albumentations.Transpose(p=0.5),\n        albumentations.HorizontalFlip(p=0.5),\n        albumentations.VerticalFlip(p=0.5),\n        albumentations.ShiftScaleRotate(p=0.5),\n        albumentations.HueSaturationValue(\n            hue_shift_limit=0.2, \n            sat_shift_limit=0.2, \n            val_shift_limit=0.2, \n            p=0.5\n        ),\n        albumentations.RandomBrightnessContrast(\n            brightness_limit=(-0.1,0.1), \n            contrast_limit=(-0.1, 0.1), \n            p=0.5\n        ),\n        albumentations.CoarseDropout(p=0.5),\n        albumentations.Cutout(p=0.5)\n    ]\n)\n\ndf_test = pd.read_csv(\"../input/cassava-leaf-disease-classification/sample_submission.csv\")\ntest_image_paths = \"../input/cassava-leaf-disease-classification/test_images/\"\n\ntest_images = df_test.image_id.values.tolist()\ntest_images = [\n    os.path.join(test_image_paths, i) for i in test_images\n]\n\ntest_dataset = ImageDataset(\n    image_paths=test_images,\n    targets=[0]*len(test_images),\n    resize=None,\n    augmentations=test_aug,\n)\n\n\ntest_loader = torch.utils.data.DataLoader(\n    test_dataset,\n    batch_size=FLAGS['batch_size'],\n    num_workers=FLAGS['num_workers'],\n    drop_last=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"final_preds = None\nfor fold in range(FLAGS['num_folds']):\n    state_dict = torch.load(os.path.join(FLAGS['model_path'],f\"xla_trained_model_{FLAGS['epochs']}_epochs_fold_{fold}.pth\"))\n    MX.load_state_dict(state_dict)\n    \n    fold_preds = tta(num_times=5, data_loader=test_loader, model=MX, device=device)\n    \n    if final_preds is None:\n        final_preds = fold_preds\n    else:\n        final_preds += fold_preds \nfinal_preds /= FLAGS['num_folds']","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"final_preds = final_preds.argmax(axis=1)\ndf_test.label = final_preds","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_test.to_csv(\"submission.csv\", index=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Now, **WE ARE DONE!**\n\nIf you enjoyed this kernel, please give it an upvote. If you have any questions or suggestions, please leave a comment!\n\nMake sure to check out the training kernel [here](https://www.kaggle.com/tanlikesmath/cassava-pytorch-xla-tpu-starter-training).\n\nAlso, check out my [related kernel](https://www.kaggle.com/tanlikesmath/the-ultimate-pytorch-tpu-tutorial-jigsaw-xlm-r) with more detailed information on PyTorch XLA/TPU training."}],"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}