{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"This Kernel is based on this amazing [⚡Plant2021 PyTorch Lightning Starter [ Training ]⚡](https://www.kaggle.com/pegasos/plant2021-pytorch-lightning-starter-training) by [Sh1r0](https://www.kaggle.com/pegasos). This kernel is intended to showcase [Weights and Biases](https://wandb.ai/site) integration with PyTorch Lightning. \n\n# ⚡ PyTorch Lightning\n\nPyTorch is an extremely powerful framework for your deep learning research. But once the research gets complicated and things like 16-bit precision, multi-GPU training, and TPU training get mixed in, users are likely to introduce bugs. **PyTorch Lightning lets you decouple research from engineering.**\n\n**PyTorch Lightning ⚡ is not another framework but a style guide for PyTorch.**\n\nTo learn more about PyTorch Lightning check out my blog posts at Weights and Biases [Fully Connected](https://wandb.ai/fully-connected):\n\n* [Image Classification using PyTorch Lightning](https://wandb.ai/wandb/wandb-lightning/reports/Image-Classification-using-PyTorch-Lightning--VmlldzoyODk1NzY)\n* [Transfer Learning Using PyTorch Lightning](https://wandb.ai/wandb/wandb-lightning/reports/Transfer-Learning-Using-PyTorch-Lightning--VmlldzoyODk2MjA)\n* [Multi-GPU Training Using PyTorch Lightning](https://wandb.ai/wandb/wandb-lightning/reports/Multi-GPU-Training-Using-PyTorch-Lightning--VmlldzozMTk3NTk)\n\n# <img src=\"https://i.imgur.com/gb6B4ig.png\" width=\"400\" alt=\"Weights & Biases\" />\n\nWeights & Biases helps you build better models faster with a central dashboard for your machine learning projects. It not only logs your training metrics but can log hyperparameters and output metrics, then visualize and compare results and quickly share findings with your team mates. Track everything you need to make your models reproducible with Weights & Biases— from hyperparameters and code to model weights and dataset versions. \n\n### [Check this Kaggle kernel to learn more about Weights and Biases$\\rightarrow$](https://www.kaggle.com/ayuraj/experiment-tracking-with-weights-and-biases)\n![img](https://i.imgur.com/BGgfZj3.png)\n\n# PyTorch Lightning + Weights and Biases \n\nPyTorch Lightning provides a lightweight wrapper for organizing your PyTorch code and easily adding advanced features such as distributed training and 16-bit precision. W&B provides a lightweight wrapper for logging your ML experiments. It is incorporated directly into the PyTorch Lightning library, so you can check out [their documentation](https://pytorch-lightning.readthedocs.io/en/stable/extensions/generated/pytorch_lightning.loggers.WandbLogger.html#pytorch_lightning.loggers.WandbLogger) for the API and reference info.\n\n### Use the intergration in few lines of code.\n\n```\nfrom pytorch_lightning.loggers import WandbLogger  # newline 1\nfrom pytorch_lightning import Trainer\n\nwandb_logger = WandbLogger()  # newline 2\ntrainer = Trainer(logger=wandb_logger)\n```\n\n[![thumbnail](https://i.imgur.com/M7xZ04g.png)](https://www.youtube.com/watch?v=hUXQm46TAKc)\n","metadata":{}},{"cell_type":"markdown","source":"# 🧰 Imports and Setups","metadata":{}},{"cell_type":"code","source":"!pip install --upgrade -q wandb\n\n# Install timm \n!pip install -q timm","metadata":{"execution":{"iopub.status.busy":"2021-12-08T17:33:32.166011Z","iopub.execute_input":"2021-12-08T17:33:32.166374Z","iopub.status.idle":"2021-12-08T17:33:50.009581Z","shell.execute_reply.started":"2021-12-08T17:33:32.166285Z","shell.execute_reply":"2021-12-08T17:33:50.008675Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Import WandbLogger\n\nCoupled with [Weights & Biases integration](https://docs.wandb.com/library/integrations/lightning), you can quickly train and monitor models for full traceability and reproducibility with only 2 extra lines of code:\n\n```python\nfrom pytorch_lightning.loggers import WandbLogger\nwandb_logger = WandbLogger()\n```\nCheck out the documentation [here](https://pytorch-lightning.readthedocs.io/en/stable/extensions/generated/pytorch_lightning.loggers.WandbLogger.html#pytorch_lightning.loggers.WandbLogger).\n","metadata":{}},{"cell_type":"code","source":"import wandb\nfrom pytorch_lightning.loggers import WandbLogger\n\nwandb.login()","metadata":{"execution":{"iopub.status.busy":"2021-12-08T17:33:50.011955Z","iopub.execute_input":"2021-12-08T17:33:50.012496Z","iopub.status.idle":"2021-12-08T17:34:42.523419Z","shell.execute_reply.started":"2021-12-08T17:33:50.012455Z","shell.execute_reply":"2021-12-08T17:34:42.522718Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import cv2\nimport timm\nimport torch\nimport numpy as np\nimport pandas as pd\n\nimport torch.nn as nn\nimport albumentations as A\nimport pytorch_lightning as pl\nimport matplotlib.pyplot as plt\n\nfrom torch.utils.data import Dataset, DataLoader\nfrom albumentations.core.composition import Compose, OneOf\nfrom albumentations.augmentations.transforms import CLAHE, GaussNoise, ISONoise\nfrom albumentations.pytorch import ToTensorV2\n\nfrom pytorch_lightning import Trainer, seed_everything\nfrom pytorch_lightning import Callback\nfrom pytorch_lightning.loggers import CSVLogger\nfrom pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping\n\nfrom sklearn.model_selection import train_test_split","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-12-08T17:34:48.267194Z","iopub.execute_input":"2021-12-08T17:34:48.267614Z","iopub.status.idle":"2021-12-08T17:34:50.334064Z","shell.execute_reply.started":"2021-12-08T17:34:48.267579Z","shell.execute_reply":"2021-12-08T17:34:50.333284Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 📀 Hyperparameters","metadata":{}},{"cell_type":"code","source":"# Config dictionary that will be logged to W&B.\nCONFIG = dict (\n    seed = 42,\n    train_val_split = 0.2,\n    model_name = 'resnet50',\n    pretrained = True,\n    img_size = 256,\n    num_classes = 12,\n    lr = 5e-4,\n    min_lr = 1e-6,\n    t_max = 20,\n    num_epochs = 10,\n    batch_size = 32,\n    accum = 1,\n    precision = 16,\n    n_fold = 5,\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n)\n\n# Directories\nPATH = \"../input/plant-pathology-2021-fgvc8/\"\n\nimage_size = CONFIG['img_size']\nTRAIN_DIR = f'../input/resized-plant2021/img_sz_{image_size}/'\nTEST_DIR = PATH + 'test_images/'\n\n# Seed everything\nseed_everything(CONFIG['seed'])","metadata":{"execution":{"iopub.status.busy":"2021-12-08T17:35:45.041658Z","iopub.execute_input":"2021-12-08T17:35:45.042004Z","iopub.status.idle":"2021-12-08T17:35:45.103409Z","shell.execute_reply.started":"2021-12-08T17:35:45.041973Z","shell.execute_reply":"2021-12-08T17:35:45.102453Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 🔧 DataModule","metadata":{}},{"cell_type":"code","source":"# Read CSV file\ndf = pd.read_csv(PATH + \"train.csv\")\n\n# Label encode \nlabels = list(df['labels'].value_counts().keys())\nlabels_dict = dict(zip(labels, range(12)))\ndf = df.replace({\"labels\": labels_dict})\ndf.head()","metadata":{"execution":{"iopub.status.busy":"2021-12-08T17:35:52.597149Z","iopub.execute_input":"2021-12-08T17:35:52.597487Z","iopub.status.idle":"2021-12-08T17:35:52.67904Z","shell.execute_reply.started":"2021-12-08T17:35:52.597454Z","shell.execute_reply":"2021-12-08T17:35:52.678018Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class PlantDataset(Dataset):\n    def __init__(self, df, transform=None):\n        self.image_id = df['image'].values\n        self.labels = df['labels'].values\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.labels)\n\n    def __getitem__(self, idx):\n        image_id = self.image_id[idx]\n        label = self.labels[idx]\n        \n        image_path = TRAIN_DIR + image_id\n        image = cv2.imread(image_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        augmented = self.transform(image=image)\n        image = augmented['image']\n        return {'image':image, 'target': label}","metadata":{"execution":{"iopub.status.busy":"2021-12-08T17:35:59.551968Z","iopub.execute_input":"2021-12-08T17:35:59.552285Z","iopub.status.idle":"2021-12-08T17:35:59.559101Z","shell.execute_reply.started":"2021-12-08T17:35:59.552254Z","shell.execute_reply":"2021-12-08T17:35:59.557913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class PlantDataModule(pl.LightningDataModule):\n    def __init__(self, batch_size, data_dir: str = './'):\n        super().__init__()\n        self.batch_size = batch_size\n        \n        # Train augmentation policy\n        self.train_transform = Compose([\n            A.RandomResizedCrop(height=CONFIG['img_size'], width=CONFIG['img_size']),\n            A.HorizontalFlip(p=0.5),\n            A.ShiftScaleRotate(p=0.5),\n            A.RandomBrightnessContrast(p=0.5),\n            A.Normalize(),\n            ToTensorV2(),\n        ])\n\n        # Validation/Test augmentation policy\n        self.test_transform = Compose([\n            A.Resize(height=CONFIG['img_size'], width=CONFIG['img_size']),\n            A.Normalize(),\n            ToTensorV2(),\n        ])\n        \n\n    def setup(self, stage=None):\n        # Assign train/val datasets for use in dataloaders\n        if stage == 'fit' or stage is None:\n            # Random train-validation split\n            train_df, valid_df = train_test_split(df, test_size=CONFIG['train_val_split'])\n            \n            # Train dataset\n            self.train_dataset = PlantDataset(train_df, self.train_transform)\n            # Validation dataset\n            self.valid_dataset = PlantDataset(valid_df, self.test_transform)\n                        \n    def train_dataloader(self):\n        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=4, drop_last=True)\n\n    def val_dataloader(self):\n        return DataLoader(self.valid_dataset, batch_size=self.batch_size, num_workers=4)","metadata":{"execution":{"iopub.status.busy":"2021-12-08T17:36:02.807659Z","iopub.execute_input":"2021-12-08T17:36:02.808025Z","iopub.status.idle":"2021-12-08T17:36:02.81998Z","shell.execute_reply.started":"2021-12-08T17:36:02.807991Z","shell.execute_reply":"2021-12-08T17:36:02.818633Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 🎺 LightningModule - Define the System","metadata":{}},{"cell_type":"code","source":"class CustomResNet(nn.Module):\n    def __init__(self, model_name='resnet18', pretrained=False):\n        super().__init__()\n        self.model = timm.create_model(model_name, pretrained=pretrained)\n        in_features = self.model.get_classifier().in_features\n        self.model.fc = nn.Linear(in_features, CONFIG['num_classes'])\n\n    def forward(self, x):\n        x = self.model(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2021-12-08T17:36:06.615871Z","iopub.execute_input":"2021-12-08T17:36:06.616194Z","iopub.status.idle":"2021-12-08T17:36:06.622051Z","shell.execute_reply.started":"2021-12-08T17:36:06.616164Z","shell.execute_reply":"2021-12-08T17:36:06.620956Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class LitCassava(pl.LightningModule):\n    def __init__(self, model):\n        super(LitCassava, self).__init__()\n        self.model = model\n        self.metric = pl.metrics.F1(num_classes=CONFIG['num_classes'])\n        self.criterion = nn.CrossEntropyLoss()\n        self.lr = CONFIG['lr']\n\n    def forward(self, x, *args, **kwargs):\n        return self.model(x)\n\n    def configure_optimizers(self):\n        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)\n        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=CONFIG['t_max'], eta_min=CONFIG['min_lr'])\n\n        return {'optimizer': self.optimizer, 'lr_scheduler': self.scheduler}\n\n    def training_step(self, batch, batch_idx):\n        image = batch['image']\n        target = batch['target']\n        output = self.model(image)\n        loss = self.criterion(output, target)\n        score = self.metric(output.argmax(1), target)\n        logs = {'train_loss': loss, 'train_f1': score, 'lr': self.optimizer.param_groups[0]['lr']}\n        self.log_dict(\n            logs,\n            on_step=False, on_epoch=True, prog_bar=True, logger=True\n        )\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        image = batch['image']\n        target = batch['target']\n        output = self.model(image)\n        loss = self.criterion(output, target)\n        score = self.metric(output.argmax(1), target)\n        logs = {'valid_loss': loss, 'valid_f1': score}\n        self.log_dict(\n            logs,\n            on_step=False, on_epoch=True, prog_bar=True, logger=True\n        )\n        return loss","metadata":{"execution":{"iopub.status.busy":"2021-12-08T17:36:09.787042Z","iopub.execute_input":"2021-12-08T17:36:09.787381Z","iopub.status.idle":"2021-12-08T17:36:09.799911Z","shell.execute_reply.started":"2021-12-08T17:36:09.787347Z","shell.execute_reply":"2021-12-08T17:36:09.798722Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 📲 Callbacks\n","metadata":{}},{"cell_type":"code","source":"# Checkpoint\ncheckpoint_callback = ModelCheckpoint(monitor='valid_loss',\n                                      save_top_k=1,\n                                      save_last=True,\n                                      save_weights_only=True,\n                                      filename='checkpoint/{epoch:02d}-{valid_loss:.4f}-{valid_f1:.4f}',\n                                      verbose=False,\n                                      mode='min')\n\n# Earlystopping\nearlystopping = EarlyStopping(monitor='valid_loss', patience=3, mode='min')","metadata":{"execution":{"iopub.status.busy":"2021-12-08T17:36:15.646801Z","iopub.execute_input":"2021-12-08T17:36:15.647141Z","iopub.status.idle":"2021-12-08T17:36:15.676013Z","shell.execute_reply.started":"2021-12-08T17:36:15.647111Z","shell.execute_reply":"2021-12-08T17:36:15.675158Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Custom Callback\nclass ImagePredictionLogger(Callback):\n    def __init__(self, val_samples, num_samples=32):\n        super().__init__()\n        self.num_samples = num_samples\n        self.val_imgs, self.val_labels = val_samples['image'], val_samples['target']\n        \n    def on_validation_epoch_end(self, trainer, pl_module):\n        # Bring the tensors to CPU\n        val_imgs = self.val_imgs.to(device=pl_module.device)\n        val_labels = self.val_labels.to(device=pl_module.device)\n        # Get model prediction\n        logits = pl_module(val_imgs)\n        preds = torch.argmax(logits, -1)\n        # Log the images as wandb Image\n        trainer.logger.experiment.log({\n            \"examples\":[wandb.Image(x, caption=f\"Pred:{pred}, Label:{y}\") \n                           for x, pred, y in zip(val_imgs[:self.num_samples], \n                                                 preds[:self.num_samples], \n                                                 val_labels[:self.num_samples])]\n            }, commit=False)","metadata":{"execution":{"iopub.status.busy":"2021-12-08T17:36:24.052387Z","iopub.execute_input":"2021-12-08T17:36:24.05282Z","iopub.status.idle":"2021-12-08T17:36:24.061664Z","shell.execute_reply.started":"2021-12-08T17:36:24.052777Z","shell.execute_reply":"2021-12-08T17:36:24.060606Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"> 📌 Tip: When logging manually through `wandb.log` or `trainer.logger.experiment.log`, make sure to use `commit=False` so the logging step does not increase.","metadata":{}},{"cell_type":"markdown","source":"## ⚡ Train and Evaluate the Model with W&B\n","metadata":{}},{"cell_type":"code","source":"# Init our data pipeline\ndatamodule = PlantDataModule(batch_size=CONFIG['batch_size'])\ndatamodule.setup()\n\n# Samples required by the custom ImagePredictionLogger callback to log image predictions.\nval_samples = next(iter(datamodule.val_dataloader()))\nval_imgs, val_labels = val_samples['image'], val_samples['target']\nval_imgs.shape, val_labels.shape","metadata":{"execution":{"iopub.status.busy":"2021-12-08T17:36:30.00176Z","iopub.execute_input":"2021-12-08T17:36:30.002104Z","iopub.status.idle":"2021-12-08T17:36:31.073511Z","shell.execute_reply.started":"2021-12-08T17:36:30.002072Z","shell.execute_reply":"2021-12-08T17:36:31.072718Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Init our model\nmodel = CustomResNet(model_name=CONFIG['model_name'], pretrained=CONFIG['pretrained'])\nlit_model = LitCassava(model)","metadata":{"execution":{"iopub.status.busy":"2021-12-08T17:36:42.84225Z","iopub.execute_input":"2021-12-08T17:36:42.842636Z","iopub.status.idle":"2021-12-08T17:36:45.619819Z","shell.execute_reply.started":"2021-12-08T17:36:42.842601Z","shell.execute_reply":"2021-12-08T17:36:45.618842Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Check out the documentation for WandbLogger [here](https://pytorch-lightning.readthedocs.io/en/stable/extensions/generated/pytorch_lightning.loggers.WandbLogger.html#pytorch_lightning.loggers.WandbLogger).\n\n> 📌 Tip: dditional arguments like entity, group, tags, etc. used by `wandb.init()` can be passed as keyword arguments in this logger.","metadata":{}},{"cell_type":"code","source":"## Initialize wandb logger\nwandb_logger = WandbLogger(project='plant-pathology-lightning', \n                           config=CONFIG,\n                           group='ResNet', \n                           job_type='train')\n\n# Initialize a trainer\ntrainer = Trainer(\n            max_epochs=CONFIG['num_epochs'],\n            gpus=1,\n            accumulate_grad_batches=CONFIG['accum'],\n            precision=CONFIG['precision'],\n            callbacks=[earlystopping,\n                       ImagePredictionLogger(val_samples)],\n            checkpoint_callback=checkpoint_callback,\n            logger=wandb_logger,\n            weights_summary='top',\n)\n\n# Train the model ⚡🚅⚡\ntrainer.fit(lit_model, datamodule)\n\n# Close wandb run\nwandb.finish() ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualize Metrics\n\n![img](https://i.imgur.com/n6P7K4M.gif)\n\n## Visualize Model Predictions\n\n![img](https://i.imgur.com/lgkLnrt.gif)\n\n## Visualize CPU and GPU Metrics\n\n![img](https://i.imgur.com/ZLjrbhj.gif)\n\n# ❄️ Resources\n\nI hope you find this kernel useful and will encouage you to try out Weights and Biases. Here are some relevant links that you might want to check out:\n\n* Check out the [official documentation](https://docs.wandb.ai/) to learn more about the best practices and advanced features. \n\n* Check out the [examples GitHub repository](https://github.com/wandb/examples) for curated and minimal examples. This can be a good starting point. \n\n* [Weights and Biases Fully Connected](https://wandb.ai/fully-connected) is a home for curated tutorials, free-form dicussions, paper summaries, industry expert advices and more. ","metadata":{}}]}