{"cells":[{"metadata":{},"cell_type":"markdown","source":"> NOTE: This notebook doesn't utilize TPU properly. TPU idle time is more than 50% most of the time."},{"metadata":{},"cell_type":"markdown","source":"## TPU Setup"},{"metadata":{"trusted":true,"_kg_hide-output":true,"collapsed":true},"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true,"collapsed":true},"cell_type":"code","source":"!pip install pytorch-lightning","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Import Packages"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import json\nimport os\nfrom typing import Dict, List\nimport logging\n\nimport numpy as np\nimport pandas as pd\n\nimport tqdm\n\nimport matplotlib.pyplot as plt\n%matplotlib inline\nimport seaborn as sns\nsns.set_style('darkgrid')\n\nimport cv2\nimport albumentations as A\nfrom albumentations.core.composition import Compose\nfrom albumentations.pytorch import ToTensorV2\n\nfrom torch.utils.data import Dataset, TensorDataset, DataLoader\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn import metrics\n\nimport pytorch_lightning as pl\nfrom pytorch_lightning import Trainer\nfrom pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping\n\nimport torch\nimport torchvision.models as models\nfrom torch import nn\nfrom torch.optim import AdamW, Adam\nimport torch.nn.functional as F\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\n\nimport torch_xla.core.xla_model as xm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Set seed for everythin(numpy, torch and python)"},{"metadata":{"trusted":true},"cell_type":"code","source":"from pytorch_lightning import seed_everything\nseed_everything(42)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Set Directories"},{"metadata":{"trusted":true},"cell_type":"code","source":"ROOT_DIR = '/kaggle/input/cassava-leaf-disease-classification/'\nTRAIN_IMAGES_FOLDER = 'train_images'\nTEST_IMAGES_FOLDER = 'test_images'\nTRAIN_CSV = 'train.csv'\nSAMPLE_SUBMISSION_CSV = 'sample_submission.csv'\nLABEL_NUM_TO_DISEASE_MAP_JSON = 'label_num_to_disease_map.json'\n\n\nimage_dir = os.path.join(ROOT_DIR, TRAIN_IMAGES_FOLDER)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Set HyperParameters"},{"metadata":{"trusted":true},"cell_type":"code","source":"LEARNING_RATE = 1e-4\nMAX_EPOCHS = 1\nBATCH_SIZE = 4","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Other Parameters"},{"metadata":{"trusted":true},"cell_type":"code","source":"NUM_WORKERS = 4\n\nIMAGE_HEIGHT = 512\nIMAGE_WIDTH = 512","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Reading Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"with open(os.path.join(ROOT_DIR, LABEL_NUM_TO_DISEASE_MAP_JSON), 'r') as file:\n    label_to_disease = json.load(file)\n    print(json.dumps(label_to_disease, indent=4))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"label_to_disease_mapping = {int(key): value for key, value in label_to_disease.items()}\nlabel_to_disease_mapping","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(os.path.join(ROOT_DIR, TRAIN_CSV))\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Preparing Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class ImageDataset(Dataset):\n    \"\"\"\n    Cassava Leaf Dataset\n    \"\"\"\n    def __init__(self,\n                image_names: List[str],\n                labels: List[int],\n                image_dir: str, \n                transforms,\n                labels_to_ohe: bool=False,\n                num_class: int = 5):        \n        self.image_names = image_names\n        self.image_dir = image_dir\n        self.transforms = transforms\n        self.num_class = num_class\n\n        if labels_to_ohe:\n            self.labels = np.zeros((len(labels), num_class))\n            self.labels[np.arange(len(labels)), np.array(labels)] = 1\n        else:\n            self.labels = np.array(labels)\n\n\n    def __len__(self):\n        return len(self.image_names)\n\n    def __getitem__(self, idx: int)->Dict[str, np.array]:\n        image_path = os.path.join(self.image_dir, self.image_names[idx])        \n        image = cv2.imread(image_path, cv2.IMREAD_COLOR)        \n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)    \n\n        target = self.labels[idx]\n\n        transformed_image = self.transforms(image=image)['image']\n        sample = {'image_path': image_path, 'image': transformed_image, 'target': torch.tensor(target)}\n\n        return sample\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"class ImageDataModule(pl.LightningDataModule):\n    def __init__(self,\n                 df,\n                 train_transforms,\n                 valid_transforms,\n                 image_dir,\n                 fold_num=0):\n        super().__init__()\n        self.df = df\n        self.train_transforms = train_transforms\n        self.valid_transforms = valid_transforms\n        self.image_dir = image_dir\n        self.fold_num = fold_num\n    \n    def prepare_data(self):\n        pass\n\n    def setup(self, stage=None):\n        \n        folds = StratifiedKFold(n_splits=5, shuffle=True)\n        \n        train_indexes, valid_indexes = list(folds.split(self.df, self.df['label']))[self.fold_num]\n        \n        train_df = self.df.iloc[train_indexes]\n        valid_df = self.df.iloc[valid_indexes]\n\n        self.train_dataset = ImageDataset(image_names=train_df.image_id.values, \n                                        labels=train_df.label.values, \n                                        image_dir=self.image_dir, \n                                        transforms=self.train_transforms)\n\n        self.valid_dataset = ImageDataset(image_names=valid_df.image_id.values, \n                                        labels=valid_df.label.values, \n                                        image_dir=self.image_dir, \n                                        transforms=self.valid_transforms)                                        \n\n    def train_dataloader(self):  \n        \"\"\"\n        sampler = torch.utils.data.distributed.DistributedSampler(\n            self.train_dataset,\n            num_replicas=xm.xrt_world_size(),\n            rank=xm.get_ordinal(),\n            shuffle=True)\n        \"\"\"    \n        train_loader = DataLoader(\n            self.train_dataset,\n            batch_size=BATCH_SIZE,\n            #sampler=sampler,\n            num_workers=NUM_WORKERS,            \n            shuffle=True\n        )\n        return train_loader\n\n    def val_dataloader(self):    \n        \"\"\"\n        sampler = torch.utils.data.distributed.DistributedSampler(\n            self.valid_dataset,\n            num_replicas=xm.xrt_world_size(),\n            rank=xm.get_ordinal(),\n            shuffle=False)        \n        \"\"\"\n        valid_loader = DataLoader(\n            self.valid_dataset,\n            batch_size=BATCH_SIZE,\n            #sampler=sampler,\n            num_workers=NUM_WORKERS,            \n            shuffle=False\n        )\n        return valid_loader\n\n    def test_dataloader(self):\n        return None\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Image Augmentation for Train and Test\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_augs = A.Compose([    \n    A.RandomResizedCrop(height=IMAGE_HEIGHT, width=IMAGE_WIDTH, p=1.0),\n    A.Flip(),    \n    A.RandomBrightnessContrast(),\n    A.ShiftScaleRotate(),\n    A.OneOf([\n            A.MotionBlur(p=.2),\n            A.MedianBlur(blur_limit=3, p=0.1),\n            A.Blur(blur_limit=3, p=0.1),\n        ], p=0.2),\n    A.Normalize(),\n    ToTensorV2(),\n])\n\nvalid_augs = A.Compose([\n    A.Resize(height=IMAGE_HEIGHT, width=IMAGE_WIDTH, p=1.0),\n    A.Normalize(),\n    ToTensorV2(),\n])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ClassifierModule(pl.LightningModule):\n    def __init__(self, learning_rate=LEARNING_RATE):\n        super().__init__()        \n        self.metric = pl.metrics.Accuracy()\n        self.learning_rate = learning_rate        \n        self.model = models.resnet101(pretrained=True)        \n        self.model.fc = nn.Linear(in_features=self.model.fc.in_features, out_features=5)                \n        \n    def forward(self, x):\n        batch_size, _, _, _ = x.shape\n        x = self.model(x)        \n        \n        return x.reshape(batch_size, -1)\n        \n\n    def configure_optimizers(self):\n        optimizer = AdamW(self.model.parameters(), lr=self.learning_rate, weight_decay=0.001)\n        scheduler = ReduceLROnPlateau(optimizer, factor=0.1, patience=2)\n\n        return (\n            [optimizer],\n            [{'scheduler': scheduler, 'interval': 'epoch', 'monitor': 'valid_loss'}],\n        )    \n    \n    def _get_loss(self, y_hat, y):\n        return nn.CrossEntropyLoss()(y_hat, y)\n\n    def training_step(self, batch, batch_idx):\n        image = batch['image']\n        y = batch['target']\n        y_hat = self(image)\n        y_hat_argmax = y_hat.argmax(1)\n        loss = self._get_loss(y_hat, y)        \n        score = self.metric(y_hat_argmax, y)        \n        \n        logs = {'train_loss': loss, 'train_accuracy': score}\n        return {\n            'loss': loss,\n            'log': logs,\n            'progress_bar': logs,\n            'logits': y_hat,\n            'target': y,\n            'train_accuracy': score,\n        }        \n    \n    def training_epoch_end(self, outputs):\n        avg_loss = torch.stack([x['loss'] for x in outputs]).mean()\n        y_true = torch.cat([x['target'] for x in outputs])\n        y_pred = torch.cat([x['logits'] for x in outputs])\n        score = self.metric(y_pred.argmax(1), y_true)\n        \n        logs = {'train_loss': avg_loss, 'train_accuracy': score}\n        \n        return {'log': logs, 'progress_bar': logs}\n\n    def validation_step(self, batch, batch_idx):\n        image = batch['image']\n        y = batch['target']\n        y_hat = self(image)\n        y_hat_argmax = y_hat.argmax(1)\n        loss = self._get_loss(y_hat, y)\n        score = self.metric(y_hat_argmax, y)\n        logs = {'valid_loss': loss, 'valid_accuracy': score}                \n\n        return {\n            'loss': loss,\n            'log': logs,\n            'progress_bar': logs,\n            'logits': y_hat,\n            'target': y,\n            f'valid_accuracy': score,\n        }        \n\n    def validation_epoch_end(self, outputs):\n        avg_loss = torch.stack([x['loss'] for x in outputs]).mean()\n        y_true = torch.cat([x['target'] for x in outputs])\n        y_pred = torch.cat([x['logits'] for x in outputs])\n        score = self.metric(y_pred.argmax(1), y_true)\n        \n        logs = {'valid_loss': avg_loss, f'valid_accuracy': score, 'accuracy': score}\n                \n        return {'valid_loss': avg_loss, 'log': logs, 'progress_bar': logs}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Data Module, change fold_num\nfold_num = 0\ndata_module = ImageDataModule(df=train_df, train_transforms=train_augs, valid_transforms=valid_augs, image_dir=image_dir, fold_num=fold_num)\n\ntrainer = pl.Trainer(\n        deterministic=True,\n        checkpoint_callback=ModelCheckpoint(monitor='train_loss', save_top_k=1, filepath='resnet101-foldnum-0_{epoch}_{valid_loss:.4f}_{accuracy:.4f}', mode='min'),\n        #gpus=1 if torch.cuda.is_available() else 0,        \n        tpu_cores=8,\n        max_epochs=MAX_EPOCHS,\n        num_sanity_val_steps=1,        \n        weights_summary='top',\n        callbacks = [EarlyStopping(monitor='valid_loss', patience=5, mode='min')]\n)\n\n\nlightning = ClassifierModule()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trainer.fit(lightning, data_module)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Resources:\n\n1 - Inspired by Artgor's notebook [Cassava disease identification with lightning](https://www.kaggle.com/artgor/cassava-disease-identification-with-lightning/notebook)\n\n2 - [TPU SUPPORT](https://pytorch-lightning.readthedocs.io/en/stable/tpu.html)"}],"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}