{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Plant 2021 with PyTorch Lightning\nThis notebook is a sample code using pytorch lightning.\nBy using pytorch lightning, you can use AMP and batch accumulation without changing the code.\n\nThe following is a notebook that uses the learned model for inference.\n[Inference notebook](https://www.kaggle.com/pegasos/plant2021-pytorch-lightning-starter-inference)"},{"metadata":{"trusted":true},"cell_type":"code","source":"package_paths = [\n    '../input/pytorch-image-library/pytorch-image-models-master/pytorch-image-models-master',\n]\nimport sys;\n\nfor pth in package_paths:\n    sys.path.append(pth)\n\nimport timm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Import"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport cv2\nimport timm\nimport torch\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 StratifiedKFold","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(f\"PyTorch Lightning version: {pl.__version__}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Config"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CFG:\n    seed = 42\n    model_name = 'tf_efficientnet_b3_ns'\n    pretrained = True\n    img_size = 512\n    num_classes = 12\n    lr = 2.5e-4\n    min_lr = 1e-6\n    t_max = 20\n    num_epochs = 20\n    batch_size = 16\n    accum = 1\n    precision = 16\n    n_fold = 5\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Load images that have been pre-resized by AnkurSingh to speed up the learning process. https://www.kaggle.com/c/plant-pathology-2021-fgvc8/discussion/227032"},{"metadata":{"trusted":true},"cell_type":"code","source":"PATH = \"../input/plant-pathology-2021-fgvc8/\"\n\n# TRAIN_DIR = PATH + 'train_images/'\nTRAIN_DIR = \"../input/resized-plant2021/img_sz_512/\"\nTEST_DIR = PATH + 'test_images/'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"seed_everything(CFG.seed)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df_all = pd.read_csv(PATH + \"train.csv\")\ndf_all.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Label Encoding"},{"metadata":{"trusted":true},"cell_type":"code","source":"labels = list(df_all['labels'].value_counts().keys())\nlabels_dict = dict(zip(labels, range(12)))\ndf_all = df_all.replace({\"labels\": labels_dict})\ndf_all.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Split Train Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"sfk = StratifiedKFold(CFG.n_fold)\nfor train_idx, valid_idx in sfk.split(df_all['image'], df_all['labels']):\n    df_train = df_all.iloc[train_idx]\n    df_valid = df_all.iloc[valid_idx]\n    break","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(f\"train size: {len(df_train)}\")\nprint(f\"valid size: {len(df_valid)}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Define Dataset"},{"metadata":{"trusted":true},"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}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_transform(phase: str):\n    if phase == 'train':\n        return Compose([\n            A.RandomResizedCrop(height=CFG.img_size, width=CFG.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    else:\n        return Compose([\n            A.Resize(height=CFG.img_size, width=CFG.img_size),\n            A.Normalize(),\n            ToTensorV2(),\n        ])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = PlantDataset(df_train, get_transform('train'))\nvalid_dataset = PlantDataset(df_valid, get_transform('valid'))\n\ntrain_loader = DataLoader(train_dataset, batch_size=CFG.batch_size, shuffle=True, num_workers=2, pin_memory=True, drop_last=True)\nvalid_loader = DataLoader(valid_dataset, batch_size=CFG.batch_size, shuffle=False, num_workers=2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Define Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CustomResNet(nn.Module):\n    def __init__(self, model_name='resnet18', pretrained=True):\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, CFG.num_classes)\n\n    def forward(self, x):\n        x = self.model(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CustomEfficientNet(nn.Module):\n    def __init__(self, model_name='tf_efficientNet_b0_ns', pretrained=True):\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.classifier = nn.Linear(in_features, CFG.num_classes)\n\n    def forward(self, x):\n        x = self.model(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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=CFG.num_classes)\n        self.criterion = nn.CrossEntropyLoss()\n        self.lr = CFG.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=CFG.t_max, eta_min=CFG.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\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = CustomEfficientNet(model_name=CFG.model_name, pretrained=CFG.pretrained)\nlit_model = LitCassava(model)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"logger = CSVLogger(save_dir='logs/', name=CFG.model_name)\nlogger.log_hyperparams(CFG.__dict__)\ncheckpoint_callback = ModelCheckpoint(monitor='valid_loss',\n                                      save_top_k=1,\n                                      save_last=True,\n                                      save_weights_only=True,\n                                      filename='{epoch:02d}-{valid_loss:.4f}-{valid_f1:.4f}',\n                                      verbose=False,\n                                      mode='min')\n\ntrainer = Trainer(\n    max_epochs=CFG.num_epochs,\n    gpus=1,\n    accumulate_grad_batches=CFG.accum,\n    precision=CFG.precision,\n    # callbacks=[EarlyStopping(monitor='valid_loss', patience=3, mode='min')],\n    checkpoint_callback=checkpoint_callback,\n    logger=logger,\n    weights_summary='top',\n)\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"trainer.fit(lit_model, train_dataloader=train_loader, val_dataloaders=valid_loader)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Result"},{"metadata":{"trusted":true},"cell_type":"code","source":"metrics = pd.read_csv(f'{trainer.logger.log_dir}/metrics.csv')\n\ntrain_acc = metrics['train_f1'].dropna().reset_index(drop=True)\nvalid_acc = metrics['valid_f1'].dropna().reset_index(drop=True)\n    \nfig = plt.figure(figsize=(7, 6))\nplt.grid(True)\nplt.plot(train_acc, color=\"r\", marker=\"o\", label='train/f1')\nplt.plot(valid_acc, color=\"b\", marker=\"x\", label='valid/f1')\nplt.ylabel('F1', fontsize=24)\nplt.xlabel('Epoch', fontsize=24)\nplt.legend(loc='lower right', fontsize=18)\nplt.savefig(f'{trainer.logger.log_dir}/f1.png')\n\ntrain_loss = metrics['train_loss'].dropna().reset_index(drop=True)\nvalid_loss = metrics['valid_loss'].dropna().reset_index(drop=True)\n\nfig = plt.figure(figsize=(7, 6))\nplt.grid(True)\nplt.plot(train_loss, color=\"r\", marker=\"o\", label='train/loss')\nplt.plot(valid_loss, color=\"b\", marker=\"x\", label='valid/loss')\nplt.ylabel('Loss', fontsize=24)\nplt.xlabel('Epoch', fontsize=24)\nplt.legend(loc='upper right', fontsize=18)\nplt.savefig(f'{trainer.logger.log_dir}/loss.png')\\\n\nlr = metrics['lr'].dropna().reset_index(drop=True)\n\nfig = plt.figure(figsize=(7, 6))\nplt.grid(True)\nplt.plot(lr, color=\"g\", marker=\"o\", label='learning rate')\nplt.ylabel('LR', fontsize=24)\nplt.xlabel('Epoch', fontsize=24)\nplt.legend(loc='upper right', fontsize=18)\nplt.savefig(f'{trainer.logger.log_dir}/lr.png')","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}