{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Pytorch Lightning Starter - PANDA competition\n\nHere I won't explain again what this competition is about. This kernel is simply to share my (local) work on how I use pytorch lightning both for this competition and in my day to day work. I hope it can serve as a useful tutorial for fellow kagglers. This is also my first public kernel :) !\n\n#### Why use Pytorch-Lightning ?\n\nI preffer code that is in general organized in small .py files with classes. For machine learning training you should be able to modify the class you want and easily add options that you switch on and off by passing the required argument to the class or its methods. Pytorch lighntning is designed to help you easily follow a pytorch based training loop and ease modifications that you may want. Want to use a new scheduler ? Then simply modify the configure_optimizer method ! The beauty of it is that it automates all the boring stuff that clogs a pure pytorch code. All these loops, .zero_grad(), .eval(), torch.save etc. are gone and handled by the framework. You just have to focus on the ML part of it.\n\nOne of the great benefit is also the ability to automate logging with tensorboard which eases comparing multiple experiments. Obviously kaggle is not very friendly with that so I suggest reproducing the code of this kernel in a local environment and use tensorboard there.\n\nYou may ask why not simply use fastai. This is now a matter of preference. Fastai automates a lot of stuff with best practices like .fit_one_cycle. But on the other hand unless you have a lot of experience with it I find it rather opaque in what is happening behind the scenes. It's a framework designed to go with doing the fastai course so that you understand the options. If like me you learnt deep learning in a more academic environment in pure pytorch or pure tensorflow then you may find fastai hard to understand without listening to the great J. Howard courses. Similarly as soon as you want to do something a bit different it can become hard to understand how to change anything. On a personal note, I'll wait for the fastai v2 course before delving into it.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Let's install it as it not in kaggle by default.\n!pip install pytorch_lightning","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Loading data\nHere I will simply create a pytorch dataset to load the tiles from the popular notebook https://www.kaggle.com/iafoss/panda-16x128x128-tiles. A dataset should simply return all the information necessary for a sample by defining the __getitem__ and __len__ magic methods.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import zipfile\nimport datetime\nimport random\nimport os\nimport pandas as pd\nimport numpy as np\nimport PIL.Image as Image\nimport torch\nimport torch.nn as nn\nimport torch.utils.data as tdata\nimport torchvision.transforms as transforms\nimport pytorch_lightning as pl\nfrom sklearn.metrics import cohen_kappa_score\nfrom sklearn.model_selection import StratifiedKFold\nfrom mish_activation import Mish","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TileDataset(tdata.Dataset):\n\n    def __init__(self, img_zip_path, dataframe, num_tiles, transform=None):\n        \"\"\"\n        img_zip: Where the images are stored\n        dataframe: The train.csv dataframe\n        num_tiles: How many tiles should the dataset return per sample\n        transform: The function to apply to the image. Usually dataaugmentation. DO NOT DO NORMALIZATION here.\n        \"\"\"\n        # Here I am using an already existing kernel output with a zipfile. \n        # I suggest extracting files as it can lead to issue with multiprocessing.\n        self.zip_img = zipfile.ZipFile(img_zip_path) \n        self.df = dataframe\n        self.num_tiles = num_tiles\n        self.img_list = self.df['image_id'].values\n        \n        self.transform = transform\n\n    def __getitem__(self, idx):\n        img_id = self.img_list[idx]\n\n        tiles = [img_id + '_' + str(i) + '.png' for i in range(0, self.num_tiles)]\n        metadata = self.df.iloc[idx]\n        image_tiles = []\n\n        for tile in tiles:\n            image = Image.open(self.zip_img.open(tile))\n\n            if self.transform is not None:\n                image = self.transform(image)\n\n            image = 1 - image\n            image = transforms.Normalize([1.0-0.90949707, 1.0-0.8188697, 1.0-0.87795304],\n                                         [0.1279171 , 0.24528177, 0.16098117])(image)\n            image_tiles.append(image)\n\n        image_tiles = torch.stack(image_tiles, dim=0)\n\n        return {'image': image_tiles, 'provider': metadata['data_provider'],\n                'isup': metadata['isup_grade'], 'gleason': metadata['gleason_score']}\n\n    def __len__(self):\n        return len(self.img_list)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"You are better off creating an instance of TileDataset in a pytorch lightning module. But here for demonstration purposes let's try our dataset class.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"transform_train = transforms.Compose([transforms.RandomHorizontalFlip(0.5),\n                                      transforms.RandomVerticalFlip(0.5),\n                                      transforms.ToTensor()])\ntrain_df = pd.read_csv('/kaggle/input/prostate-cancer-grade-assessment/train.csv')\ntrainset = TileDataset('/kaggle/input/panda-16x128x128-tiles/train.zip', train_df, 12, transform=transform_train)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image = trainset[0]['image']\nprint(image.shape, image.mean(), image.std())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"One element of the dataset is 12 tiles, of 128x128 RGB color images in a N, C, H, W tensor.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"## Making a Model\n\nInspired by other popular kernel here is a simple resnext50 pytorch model. As explained at the start I suggest you define the head as a module itself. This will allow you to easily create new heads and then switch them in the main module.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class AdaptiveConcatPool2d(nn.Module):\n    # This layer will concatenate both average and max pool\n    def __init__(self):\n        super().__init__()\n        self.avg = nn.AdaptiveAvgPool2d(output_size=(1, 1))\n        self.max = nn.AdaptiveMaxPool2d(output_size=(1, 1))\n\n    def forward(self, x):\n        avg_x = self.avg(x)\n        max_x = self.max(x)\n        return torch.cat([avg_x, max_x], dim=1)\n\n\nclass Flatten(nn.Module):\n    def forward(self, x):\n        return x.view(x.shape[0], -1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class BasicHead(nn.Module):\n    # The head of our model\n    def __init__(self, c_in, c_out, n_tiles):\n        self.n_tiles = n_tiles\n        super().__init__()\n        self.fc = nn.Sequential(AdaptiveConcatPool2d(),\n                                Flatten(),\n                                nn.Dropout(0.5),\n                                nn.Linear(c_in * 2, 512),\n                                Mish(),\n                                nn.BatchNorm1d(512),\n                                nn.Dropout(0.5),\n                                nn.Linear(512, c_out))\n\n    def forward(self, x):\n\n        bn, c, height, width = x.shape\n        h = x.view(-1, self.n_tiles, c, height, width).permute(0, 2, 1, 3, 4) \\\n            .contiguous().view(-1, c, height * self.n_tiles, width)\n        h = self.fc(h)\n        return h","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Model(nn.Module):\n    # The mnain model combining a backbone and a head\n    def __init__(self, c_out=6, n_tiles=12, tile_size=128, backbone='resnext50_semi', head='basic', **kwargs):\n        super().__init__()\n        if backbone == 'resnext50_semi':\n            m = torch.hub.load('facebookresearch/semi-supervised-ImageNet1K-models', 'resnext50_32x4d_ssl')\n        elif backbone == 'resnet50':\n            m = models.resnet50(pretrained=True)\n        \n        c_feature = list(m.children())[-1].in_features  \n        self.feature_extractor = nn.Sequential(*list(m.children())[:-2])  # Remove resnet head\n        self.n_tiles = n_tiles\n        self.tile_size = tile_size\n        if head == 'basic':\n            self.head = BasicHead(c_feature, c_out, n_tiles)\n\n    def forward(self, x):\n        h = x.view(-1, 3, self.tile_size, self.tile_size)\n        h = self.feature_extractor(h)\n        h = self.head(h)\n\n        return h","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Pytorch Lightning module definition\n\nIn a normal pytorch code you probably would instantiate the model, dataloaders and make a nested for loop for epochs and batches. Pytorch lightning automates the engineering parts like the loops so that you focus on the ML part. To do that you create a pytorch lightning model and then define every ML step inside of it. To help you understand I have added comments under every method you need to implement.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class LightModel(pl.LightningModule):\n\n    def __init__(self, df_train, train_idx, val_idx, hparams):\n        # This is where paths and options should be stored. I also store the\n        # train_idx, val_idx for cross validation since the dataset are defined \n        # in the module !\n        super().__init__()\n        self.train_idx = train_idx\n        self.val_idx = val_idx\n        self.df_train = df_train\n\n        self.model = Model(c_out=hparams.c_out,  # This would be different for regression or classification\n                           n_tiles=hparams.n_tiles,\n                           tile_size=hparams.tile_size,\n                           backbone=hparams.backbone,\n                           head=hparams.head)\n\n        self.hparams = hparams\n        self.trainset = None\n        self.valset = None\n\n    def forward(self, batch):\n        # What to do with a batch in a forward. Usually simple if everything is already defined in the model.\n        return self.model(batch['image'])\n\n    def prepare_data(self):\n        # This is called at the start of training and is where everything data related should be initialized.\n        transform_train = transforms.Compose([transforms.RandomHorizontalFlip(0.5),\n                                              transforms.RandomVerticalFlip(0.5),\n                                              transforms.ToTensor()])\n        transform_test = transforms.Compose([transforms.ToTensor()])\n        \n        self.trainset = TileDataset('/kaggle/input/panda-16x128x128-tiles/train.zip', self.df_train.iloc[self.train_idx], self.hparams.n_tiles, transform=transform_train)\n        self.valset = TileDataset('/kaggle/input/panda-16x128x128-tiles/train.zip', self.df_train.iloc[self.val_idx], self.hparams.n_tiles, transform=transform_test)\n\n    def train_dataloader(self):\n        # Simply define a pytorch dataloader here that will take care of batching. Note it works well with dictionnaries !\n        train_dl = tdata.DataLoader(self.trainset, batch_size=BATCH_SIZE, shuffle=True,\n                                    num_workers=0)  # Using only one worker can be slow but zipfile can lead to bugs. You may try with multiple workers. In general use extracted files.\n        return train_dl\n\n    def val_dataloader(self):\n        # Same but for validation. Pytorch lightning allows multiple validation dataloaders hence why I return a list.\n        val_dl = tdata.DataLoader(self.valset, batch_size=BATCH_SIZE, shuffle=False,\n                                  num_workers=0)\n        return [val_dl]\n\n    def cross_entropy_loss(self, logits, gt):\n        # How to calculate the loss. Note this method is actually not a part of pytorch lightning ! It's only good practice\n        loss_fn = nn.CrossEntropyLoss()\n        return loss_fn(logits, gt)\n\n    def configure_optimizers(self):\n        # Optimizers and schedulers. Note that each are in lists of equal length to allow multiple optimizers (for GAN for example)\n        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.hparams.lr, weight_decay=3e-6)\n        scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=10 * self.hparams.lr, \n                                                        epochs=self.hparams.epochs, steps_per_epoch=len(self.train_dataloader()))\n        return [optimizer], [scheduler]\n\n    def training_step(self, batch, batch_idx):\n        # This is where you must define what happens during a training step (per batch)\n        logits = self(batch)\n        loss = self.cross_entropy_loss(logits, batch['isup']).unsqueeze(0)  # You need to unsqueeze in case you do multi-gpu training\n        preds = logits.argmax(1)\n        # Pytorch lightning will call .backward on what is called 'loss' in output\n        # 'log' is reserved for tensorboard and will log everything define in the dictionary\n        return {'loss': loss, 'log': {'train_loss': loss}}\n\n    def validation_step(self, batch, batch_idx):\n        # This is where you must define what happens during a validation step (per batch)\n        logits = self(batch)\n        loss = self.cross_entropy_loss(logits, batch['isup']).unsqueeze(0)\n        preds = logits.argmax(1)\n        return {'val_loss': loss, 'preds': preds, 'gt': batch['isup']}\n\n    def validation_epoch_end(self, outputs):\n        # This is what happens at the end of validation epoch. Usually gathering all predictions\n        # outputs is a list of dictionary from each step.\n        avg_loss = torch.cat([out['val_loss'] for out in outputs], dim=0).mean()\n        preds = torch.cat([out['preds'] for out in outputs], dim=0)\n        gt = torch.cat([out['gt'] for out in outputs], dim=0)\n        preds = preds.detach().cpu().numpy()\n        gt = gt.detach().cpu().numpy()\n\n        kappa = cohen_kappa_score(preds, gt, weights='quadratic')\n        tensorboard_logs = {'val_loss': avg_loss, 'kappa': kappa}\n        print(f'Epoch {self.current_epoch}: {avg_loss:.2f}, kappa: {kappa:.4f}')\n\n        return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"This is a small helper to make hparams act like args. Making it easier to then switch the code to argparser.\nThis is also what allows you to use hparams.arg instead of hparams['arg']","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import argparse\ndef dict_to_args(d):\n    args = argparse.Namespace()\n\n    def dict_to_args_recursive(args, d, prefix=''):\n        for k, v in d.items():\n            if type(v) == dict:\n                dict_to_args_recursive(args, v, prefix=k)\n            elif type(v) in [tuple, list]:\n                continue\n            else:\n                if prefix:\n                    args.__setattr__(prefix + '_' + k, v)\n                else:\n                    args.__setattr__(k, v)\n\n    dict_to_args_recursive(args, d)\n    return args\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training\nLet's start by specifying parameters, the seed and output folder.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"SEED = 33\nBATCH_SIZE = 4\nNAME = 'resnext50'\nOUTPUT_DIR = './lightning_logs'\nrandom.seed(SEED)\nos.environ['PYTHONHASHSEED'] = str(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\ntorch.cuda.manual_seed(SEED)\ntorch.backends.cudnn.deterministic = True\ntorch.backends.cudnn.benchmark = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# This is where you should specify for an experiment what the model should use.\nhparams = {'backbone': 'resnext50_semi', \n           'head': 'basic',\n           'lr': 1e-4,\n           'n_tiles': 12,\n           'c_out': 6,\n           'epochs':2,  # You obviously want to increase this :)\n           'tile_size': 128}\n\nhparams = dict_to_args(hparams) ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Before start of training, the kernel output that I am using https://www.kaggle.com/iafoss/panda-16x128x128-tiles does not process all images in the zip file (missing those with a mask). Therefore I manually remove them from the dataframe to avoid issues with the datasets. **This is not necessary if you processed all tiles !**.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv('/kaggle/input/prostate-cancer-grade-assessment/train.csv')\nzip_img = zipfile.ZipFile('/kaggle/input/panda-16x128x128-tiles/train.zip')\nnot_in_image_zip = []\nfor img_id in train_df['image_id']:\n    try:\n        zip_img.open(img_id + '_' + '0.png')\n    except KeyError:\n        not_in_image_zip.append(img_id)\ntrain_df = train_df[~train_df['image_id'].isin(not_in_image_zip)]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"For training we just need to instantiate the pytorch lightning module and a trainer with a few options. Most importantly this is where you specify how many GPU to use (or TPU) and if you want to do mixed precision training (with apex). For the purpose of this kernel I just do FP32 1GPU training but please read the pytorch lightning doc if you want to try TPU and/or mixed precision.\n\nNote: For some reason the kernel has a bit of trouble with the zipfile and I have to usually","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"kfold = StratifiedKFold(n_splits=5, random_state=SEED, shuffle=True)\nsplits = kfold.split(train_df, train_df['isup_grade'])\ndate = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')\n\n# Checkpoints and logs will be in ./OUTPUT_DIR/NAME-DATE/fold_i\nNAME = 'resnext50'\nOUTPUT_DIR = './lightning_logs'\n\nfor fold, (train_idx, val_idx) in enumerate(splits):\n    print(f'Fold {fold + 1}')\n    # Defining clearly to the tensorboard logger in order to put every fold under the same directory.\n    tb_logger = pl.loggers.TensorBoardLogger(save_dir=OUTPUT_DIR,\n                                             name=f'{NAME}' + '-' + date,\n                                             version=f'fold_{fold + 1}')\n\n    # Define what metric the checkpoint should track (can be anything returned from the validation_end method)\n    checkpoint_callback = pl.callbacks.ModelCheckpoint(filepath=tb_logger.log_dir + \"/{epoch:02d}-{kappa:.4f}\",\n                                                       monitor='kappa', mode='max')\n\n    # Initiate model\n    model = LightModel(train_df, train_idx, val_idx, hparams)\n    \n    # Define trainer\n    # Here you can \n    trainer = pl.Trainer(gpus=[0], max_nb_epochs=hparams.epochs, auto_lr_find=False,\n                         gradient_clip_val=1,\n                         logger=tb_logger,\n                         accumulate_grad_batches=1,              # BatchNorm ?\n                         checkpoint_callback=checkpoint_callback\n                         )\n    # lr_finder = trainer.lr_find(model)\n    # fig = lr_finder.plot(suggest=True)\n    # fig.savefig('lr_plot.png')\n    trainer.fit(model)\n    \n    # One last thing. In checkpoints, pytorch lightning will save the weights and the state of the optimizer. \n    # This makes the weights very large. If you want to isolate the model weight simply use:\n    torch.save(model.model.state_dict(), OUTPUT_DIR + '/' + NAME + '-' + date + '/' + f'fold_{fold}.pth')\n    # One fold training (Remove this for training all folds)\n    break","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"That's it ! Hope it was useful. ","execution_count":null}],"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}