{"cells":[{"metadata":{},"cell_type":"markdown","source":"<img src=\"https://i.imgur.com/pNKgZgL.png\" alt=\"Fastai2 and Weights & Biases\" width=\"400\"/>\n\n<div><img /></div>\n\n<img src=\"https://i.imgur.com/uEtWSEb.png\" width=\"650\" alt=\"Weights & Biases\" />\n\n<div><img /></div>\n<br>\n<br>\n\n## 💨 Fastai + Weights & Biases + Albumentations + pytorch-image-models"},{"metadata":{},"cell_type":"markdown","source":"> This notebook builds upon all the other excellent notebooks show regarding this competition. In this notebook i will show how to train a `efficientnet` model from `pytorch-image-models` using the `fastai` libraray. Additionaly we will also go over how to perform image augmentations using `albumentations` and track progress using `wandb`."},{"metadata":{},"cell_type":"markdown","source":"### What this notebook covers ?\n- model training.\n- data augmentation.\n- custom fastai `datablock`, `callbacks`\n- track model progress using `wandb`."},{"metadata":{},"cell_type":"markdown","source":"### What this notebook doesn't cover?\n\n- EDA\n- creating StratifiedKFolds"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!git clone https://github.com/benihime91/leaf-disease-classification-kaggle\n!pip install timm wandb --upgrade --quiet","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\nsys.path.append(\"./leaf-disease-classification-kaggle/\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"!wandb login a74f67fd5fae293e301ea8b6710ee0241f595a63 # [YOUR WANDB API KEY HERE]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport uuid \nimport os\nimport timm\nimport wandb\nimport pandas as pd\nimport albumentations as A\n\nfrom fastai.vision.all import *\nfrom fastai.callback.wandb import *\n\nfrom torch.distributions.beta import Beta\n\nfrom src.model import _cut_model, _num_feats, _create_head\n\nidx = uuid.uuid1()\nidx = str(idx).split(\"-\")[0]\n\nset_seed(42)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Create the training configuration"},{"metadata":{},"cell_type":"markdown","source":"Specify the config file. Here we will specify all our training config parameters like : `num_classes`, `input_dims`, paths, etc etc."},{"metadata":{"trusted":true},"cell_type":"code","source":"class Config:\n    fold_num    = 0\n    num_classes = 5\n    csv_dir     = \"./leaf-disease-classification-kaggle/data/fold_df.csv\"\n    image_dir   = \"/kaggle/input/cassava-leaf-disease-classification/train_images/\"\n    input_dims  = 256\n    model_arch  = \"efficientnet_b3a\"\n    project     = \"kaggle-leaf-disease-fastai-runs\" # [YOUR WANDB PROJECT NAME]\n    \n# init config\ncfg = Config()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prepare the dataframe with folds"},{"metadata":{},"cell_type":"markdown","source":"We already have a stratified K-Fold data in csv format. Let's load in the the data in pandas dataframe and we will modify the dataframe for this particular task ."},{"metadata":{"trusted":true},"cell_type":"code","source":"idx2lbl = {0:\"Cassava Bacterial Blight (CBB)\",\n          1:\"Cassava Brown Streak Disease (CBSD)\",\n          2:\"Cassava Green Mottle (CGM)\",\n          3:\"Cassava Mosaic Disease (CMD)\",\n          4:\"Healthy\"}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data             = pd.read_csv(cfg.csv_dir)\ndata[\"filePath\"] = [os.path.join(cfg.image_dir, data[\"image_id\"][idx]) for idx in range(len(data))]\ndata[\"is_valid\"] = [data.kfold[n] == cfg.fold_num for n in range(len(data))]\ndata['label'].replace(idx2lbl, inplace=True)\n\n# shuffle the dataset\ndata = data.sample(frac=1).reset_index(drop=True, inplace=False)\ndata.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Create custom fast.ai data block"},{"metadata":{},"cell_type":"markdown","source":"We are going to gather our data using the fast.ai `DataBlock` API."},{"metadata":{},"cell_type":"markdown","source":"Albumentations + fast.ai datablock, For data augmentation we are going to us the `albumentations` library."},{"metadata":{"trusted":true},"cell_type":"code","source":"# class for albumentations transformations\n# taken from : https://forums.fast.ai/t/albumentation-transformations-for-train-and-test-dataset/82642\n\nclass AlbumentationsTransform(RandTransform):\n    split_idx,order=None,2\n    def __init__(self, train_aug, valid_aug): store_attr()\n    \n    def before_call(self, b, split_idx):\n        self.idx = split_idx\n    \n    def encodes(self, img: PILImage):\n        if self.idx == 0:\n            aug_img = self.train_aug(image=np.array(img))['image']\n        else:\n            aug_img = self.valid_aug(image=np.array(img))['image']\n        return PILImage.create(aug_img)\n    \ndef get_train_aug(): \n    return A.Compose([\n        A.OneOf([\n            A.RandomResizedCrop(cfg.input_dims, cfg.input_dims), \n            A.CenterCrop(cfg.input_dims, cfg.input_dims)\n        ], p=0.5),\n        A.Resize(cfg.input_dims, cfg.input_dims, p=1.0),\n        A.HorizontalFlip(),\n        A.ShiftScaleRotate(),\n        A.OneOf([A.Flip(), A.Transpose(), A.IAAPerspective()]),\n        A.RandomBrightnessContrast(0.1, 0.1, p=0.5),\n        A.OneOf([A.CLAHE(), A.HueSaturationValue(0.2, 0.2, 0.2, p=0.5),], p=0.4),\n        A.OneOf([A.CoarseDropout(), A.Cutout(), A.JpegCompression()], p=0.5),\n    ])\n\ndef get_valid_aug():\n    return A.Compose([A.Resize(cfg.input_dims, cfg.input_dims, p=1.0)], p=1.)\n\nitem_tfms = AlbumentationsTransform(get_train_aug(), get_valid_aug())\n\nbatch_tfms = [Normalize.from_stats(*imagenet_stats)]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dblock = DataBlock(blocks=(ImageBlock, CategoryBlock),\n                splitter=ColSplitter(),\n                get_x=lambda o: o[\"filePath\"],\n                get_y=lambda o: o[\"label\"],\n                item_tfms=item_tfms,\n                batch_tfms=batch_tfms)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We can now create our `dataloaders` and have a look at our data"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Load the data + preprocessing & data augmentation using the fast.ai\n# datablock created above\ndls = dblock.dataloaders(data)\n\n# Let's look at a batch of data to make sure everything looks alright:\ndls.show_batch(figsize=(12,12))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Init, modify, load model for fast.ai Learner class"},{"metadata":{},"cell_type":"markdown","source":"In this section we will load a model from the `timm` library and then modify the head of the `model` for our particular. Our custom head is based on the fast.ai custom head for used in fast.ai `cnn_learner`. Our head will be \n```\nSequential(\n  (0): AdaptiveConcatPool2d(\n    (ap): AdaptiveAvgPool2d(output_size=1)\n    (mp): AdaptiveMaxPool2d(output_size=1)\n  )\n  (1): Flatten()\n  (2): BatchNorm1d( ... )\n  (3): Dropout( ... )\n  (4): Linear( ... )\n  (5): ReLU( ... )\n  (6): BatchNorm1d( ... )\n  (7): Dropout( ... )\n  (8): Linear( ... )\n)\n```\nWe will modify the `parameters` for the model for our current. For this we will use the convenience functions `_cut_model`, `_num_feats`, `_create_head`."},{"metadata":{"trusted":true},"cell_type":"code","source":"class TransferLearningModel(Module):\n    \"\"\"Transfer Learning with pre-trained encoder.\n    Args:\n        encoder    : a pre-trained model architecture like tf_efficientnet_b3\n        num_classes: total number of output classes\n        lin_ftrs   : number of Linear nodes before the final output layer\n    \"\"\"\n    def __init__(self, encoder, num_classes = cfg.num_classes, lin_ftrs=512, cut=-2):\n        # remove the classifier from the encoder which are\n        # the final two layers of the encoder\n        self.encoder = _cut_model(encoder, cut)\n        \n        # calculate output features of the cut encoder\n        ftrs = _num_feats(self.encoder) * 2\n        \n        # create the head/decoder for the encoder\n        self.decoder = _create_head(ftrs, num_classes, lin_ftrs)\n        apply_init(self.decoder)\n        \n    def forward(self, xb):\n        # 1. Feature extraction:\n        feats = self.encoder(xb)\n        \n        # 2. Decoder (returns logits):\n        logits = self.decoder(feats)\n        return logits","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We will now create custom `splitter` to pass along to `Learner` to take advantage of transfer learning. In order to use transfer learning efficiently, we will want to freeze the pretrained model at first, and train only the head. The `params` function is useful to return all parameters of the model, so we can create a simple splitter like so:"},{"metadata":{"trusted":true},"cell_type":"code","source":"def splitter(net): \n    \"\"\"\n    returns parameters of the classifer and the base of the model\n    required for gradual freezing/unfrezing and discriminative\n    lr techniques for fast.ai\n    \"\"\"\n    return [params(net.encoder), params(net.decoder)]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Init fast.ai Learner and start training"},{"metadata":{},"cell_type":"markdown","source":"Let's now create our `Learner` object."},{"metadata":{},"cell_type":"markdown","source":"Implement cutmix callback from [here](https://github.com/fastai/fastai/blob/master/fastai/callback/cutmix.py#L10) ."},{"metadata":{"trusted":true},"cell_type":"code","source":"# straight copy from : https://docs.fast.ai/callback.cutmix#CutMix\nclass CutMix(Callback):\n    \"Implementation of `https://arxiv.org/abs/1905.04899`\"\n    run_after,run_valid = [Normalize],False\n    def __init__(self, alpha=1.): self.distrib = Beta(tensor(alpha), tensor(alpha))\n    def before_fit(self):\n        self.stack_y = getattr(self.learn.loss_func, 'y_int', False)\n        if self.stack_y: self.old_lf,self.learn.loss_func = self.learn.loss_func,self.lf\n\n    def after_fit(self):\n        if self.stack_y: self.learn.loss_func = self.old_lf\n\n    def before_batch(self):\n        W, H = self.xb[0].size(3), self.xb[0].size(2)\n        lam = self.distrib.sample((1,)).squeeze().to(self.x.device)\n        lam = torch.stack([lam, 1-lam])\n        self.lam = lam.max()\n        shuffle = torch.randperm(self.y.size(0)).to(self.x.device)\n        xb1,self.yb1 = tuple(L(self.xb).itemgot(shuffle)),tuple(L(self.yb).itemgot(shuffle))\n        nx_dims = len(self.x.size())\n        x1, y1, x2, y2 = self.rand_bbox(W, H, self.lam)\n        self.learn.xb[0][:, :, x1:x2, y1:y2] = xb1[0][:, :, x1:x2, y1:y2]\n        self.lam = (1 - ((x2-x1)*(y2-y1))/float(W*H)).item()\n\n        if not self.stack_y:\n            ny_dims = len(self.y.size())\n            self.learn.yb = tuple(L(self.yb1,self.yb).map_zip(torch.lerp,weight=unsqueeze(self.lam, n=ny_dims-1)))\n\n    def lf(self, pred, *yb):\n        if not self.training: return self.old_lf(pred, *yb)\n        with NoneReduce(self.old_lf) as lf:\n            loss = torch.lerp(lf(pred,*self.yb1), lf(pred,*yb), self.lam)\n        return reduce_loss(loss, getattr(self.old_lf, 'reduction', 'mean'))\n\n    def rand_bbox(self, W, H, lam):\n        cut_rat = torch.sqrt(1. - lam)\n        cut_w = (W * cut_rat).type(torch.long)\n        cut_h = (H * cut_rat).type(torch.long)\n        # uniform\n        cx = torch.randint(0, W, (1,)).to(self.x.device)\n        cy = torch.randint(0, H, (1,)).to(self.x.device)\n        x1 = torch.clamp(cx - cut_w // 2, 0, W)\n        y1 = torch.clamp(cy - cut_h // 2, 0, H)\n        x2 = torch.clamp(cx + cut_w // 2, 0, W)\n        y2 = torch.clamp(cy + cut_h // 2, 0, H)\n        return x1, y1, x2, y2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"encoder = timm.create_model(cfg.model_arch, pretrained=True, num_classes=dls.c)\n\nmodel   = TransferLearningModel(encoder, cut=-4)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"_kg_hide-input":true,"trusted":true,"collapsed":true},"cell_type":"code","source":"model","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true,"collapsed":true},"cell_type":"code","source":"run_name = f\"{cfg.model_arch}-fold={cfg.fold_num}-{idx}\"\nwandb.init(project=cfg.project, name=run_name)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"loss  = CrossEntropyLossFlat()\nlearn = (Learner(dls, model, loss_func=loss, splitter=splitter, metrics=accuracy, cbs=[WandbCallback(), CutMix()]).to_native_fp16())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We are now provided with a Learner object. We will first freeze the model i.e., only the weights of the head can be updated. \nIn order to train a model, we need to find the most optimal learning rate, which can be done with fastai's learning rate finder:"},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.freeze()\nlearn.lr_find()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.fit_one_cycle(22, 1e-03)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":">Tip: Save model weights after each each run to note lose progress."},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(learn.model.state_dict(), \"stage-1.pt\")","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"learn.model.load_state_dict(torch.load(\"stage-1.pt\"))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Unfreeze the whole model and train"},{"metadata":{"trusted":true},"cell_type":"code","source":"# train all the parameters of the model\nlearn.unfreeze()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.lr_find()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"learn.fit_one_cycle(12, slice(1e-06, 1e-04))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We can now `save` the weights of our PyTorch model or export the fastai `Learner`. I prefere saving the weights of the model so than i can simple use one the inference notebooks already uploaded to make inference"},{"metadata":{"trusted":true},"cell_type":"code","source":"save_dir = f\"/kaggle/working/{cfg.model_arch}-{cfg.input_dims}-fold={cfg.fold_num}.pt\"\ntorch.save(learn.model.state_dict(), save_dir)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> Optional: Update saved weights file to `wandb` and finish the `wandb` `run`."},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"wandb.save(save_dir)\nwandb.finish()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## References"},{"metadata":{},"cell_type":"markdown","source":"- https://www.kaggle.com/tanlikesmath/cassava-classification-eda-fastai-starter\n- https://www.kaggle.com/khyeh0719/pytorch-efficientnet-baseline-train-amp-aug\n- https://www.kaggle.com/muellerzr/cassava-fastai-starter\n- https://www.kaggle.com/muellerzr/recreating-abhishek-s-tez-with-fastai"}],"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}