{"cells":[{"metadata":{},"cell_type":"markdown","source":"Here is an attempt to do KFold, Parallel training. However training happens only on one core at a time. "},{"metadata":{"trusted":true,"scrolled":false},"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\n!pip uninstall -q typing --yes","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"!pip install pytorch-lightning\n!pip install timm","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"scrolled":true},"cell_type":"code","source":"import os\nimport cv2\nimport pandas as pd\nimport numpy as np\nimport random\nimport pytorch_lightning as pl\nfrom pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping\nimport torch\n# from torchvision import models\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as albu\nfrom albumentations.pytorch.transforms import ToTensorV2\nfrom sklearn.model_selection import StratifiedKFold\n# from efficientnet_pytorch import EfficientNet\nimport torch_xla.core.xla_model as xm\nimport torch_xla\nimport timm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"import random\ndef seed_everything(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nseed_everything()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"TRAIN_CSV = \"../input/cassava-leaf-disease-classification/train.csv\"\nTRAIN_IMAGE_FOLDER = '../input/cassava-leaf-disease-classification/train_images'\nCLASSES = 5","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Hyper parameters"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"FOLDS = 8\nBATCH_SIZE = 8\nLR = 0.01\nEPOCHS=2\nLOSS_FUNCTION = nn.BCEWithLogitsLoss()\nIMG_SIZE = 128\nEARLY_STOPPING = True\nMODEL_ARCH = 'resnet50'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Dataset"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"class CassavaDataset(Dataset):\n    def __init__(self, train, train_mode=True, transforms=None):\n        self.train = train\n        self.transforms = transforms\n        self.train_mode = train_mode\n    \n    def __len__(self):\n        return self.train.shape[0]\n    \n    def __getitem__(self, index):\n        image_path = os.path.join(TRAIN_IMAGE_FOLDER, self.train.iloc[index].image_id)\n        image = cv2.imread(image_path, cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        if (self.transforms):\n            image = self.transforms(image=image)[\"image\"]\n        \n        if not(self.train_mode):\n            return {\"x\":image}\n        \n        return {\n            \"x\": image,\n            \"y\": torch.tensor(self.train.iloc[index, self.train.columns.str.startswith('label')], dtype=torch.float64)\n        }","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Transforms"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"def get_augmentations():\n    \n    mean = (0.485, 0.456, 0.406)\n    std = (0.229, 0.224, 0.225)    \n    \n    train_augmentations = albu.Compose([\n        albu.RandomResizedCrop(IMG_SIZE, IMG_SIZE, p=1.0),\n        albu.Transpose(p=0.5),\n        albu.HorizontalFlip(p=0.5),\n        albu.VerticalFlip(0.5),\n        albu.Normalize(mean, std, max_pixel_value=255, always_apply=True),        \n        ToTensorV2(p=1.0)\n    ], p=1.0)\n    \n    valid_augmentations = albu.Compose([\n        albu.Normalize(mean, std, max_pixel_value=255, always_apply=True),        \n        ToTensorV2(p=1.0)\n    ], p=1.0)   \n    \n    return train_augmentations, valid_augmentations\n\ntrain_augs, val_augs = get_augmentations()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### NN Model"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"# # These are the available model architectures in timm\n# from pprint import pprint\n# model_names = timm.list_models(pretrained=True)\n# pprint(model_names)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"class Model(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.model = timm.create_model(MODEL_ARCH, pretrained=True)\n#         self.model = base_model\n\n#         # Efficientnets\n#         n_features = self.model.classifier.in_features\n#         self.model.classifier = nn.Linear(n_features, CLASSES)\n        \n        # Resnets\n        n_features = self.model.fc.in_features\n        self.model.fc = nn.Linear(n_features, CLASSES)\n        \n    def forward(self, x):\n        x = self.model(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### K-Fold CV"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"traincsv = pd.read_csv(TRAIN_CSV)\ntraincsv['kfold'] = -1\ntraincsv = traincsv.sample(frac=1).reset_index(drop=True)\nstratifier = StratifiedKFold(n_splits=FOLDS)\n\nfor fold, (train_index, val_index) in enumerate(stratifier.split(X=traincsv.image_id.values, y=traincsv.label.values)):\n    traincsv.loc[val_index, \"kfold\"] = fold\n\ntraincsv.to_csv(\"train_folds.csv\", index=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### PL Data module"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"class CassavaDataModule(pl.LightningDataModule):\n    def __init__(self, fold):\n        super().__init__()\n        self.train_aug, self.valid_aug = get_augmentations()\n        self.fold = fold\n        self.batch_size = BATCH_SIZE\n    \n    def setup(self, stage=None):\n        folds = pd.read_csv('./train_folds.csv')\n        folds = pd.get_dummies(folds, columns=['label'])\n        train_fold = folds.loc[folds[\"kfold\"] != self.fold]\n        val_fold = folds.loc[folds[\"kfold\"] == self.fold]\n        \n        self.train_ds = CassavaDataset(train_fold, transforms=train_augs)\n        self.val_ds = CassavaDataset(val_fold, transforms=val_augs)\n        \n    def train_dataloader(self):\n        return DataLoader(self.train_ds, self.batch_size, num_workers=4, shuffle=True)\n        \n    def val_dataloader(self):\n        return DataLoader(self.val_ds, self.batch_size, num_workers=4, shuffle=False)        \n        ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Callbacks"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"early_stopping = EarlyStopping('val_accuracy', patience=3, mode='max')\n\ncallbacks=[]\n\nif EARLY_STOPPING == True:\n    callbacks.append(early_stopping)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### PL Module"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"class CassavaPLModule(pl.LightningModule):\n    def __init__(self, hparams, model):\n        super(CassavaPLModule, self).__init__()\n        self.hparams = hparams\n        self.model = model\n        self.criterion = LOSS_FUNCTION\n        self.accuracy = pl.metrics.Accuracy()\n        \n    def forward(self, x):\n        return self.model(x)\n    \n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.hparams.lr)\n        scheduler = {\n            'scheduler': \n                torch.optim.lr_scheduler.ReduceLROnPlateau(\n                    optimizer, patience=3,\n                    threshold=0.001,\n                    mode='min', verbose=True\n                ),\n            'interval': 'epoch',\n            'monitor' : 'val_loss'\n        }\n        return [optimizer], [scheduler]\n        \n    def training_step(self, batch, batch_index):\n        # One batch at a time\n        features = batch['x']\n        targets = batch['y']\n        out = self(features)\n        loss = self.criterion(out, targets)\n        self.log(\"train_loss\", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)      \n        metric_acc = self.accuracy(out, targets)\n        self.log(\"train_accuracy\", metric_acc, on_step=True, on_epoch=True, prog_bar=True,logger=True)\n        \n    def validation_step(self, batch, batch_index):\n        # One batch at a time\n        features = batch['x']\n        targets = batch['y']\n        out = self(features)\n        loss = self.criterion(out, targets)\n        self.log(\"val_loss\", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True) \n        metric_acc = self.accuracy(out, targets)\n        self.log(\"val_accuracy\", metric_acc, on_step=True, on_epoch=True, prog_bar=True, logger=True)        ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Training"},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"def train(fold):   \n    checkpoint_callback = ModelCheckpoint(\n        dirpath='checkpoints/',\n        filename='model_{fold}-{val_loss:.2f}',\n        monitor='val_loss', verbose=True,\n        save_last=False, save_top_k=1, save_weights_only=False,\n        mode='min', period=1, prefix=''\n    )        \n    \n    tpu_core = fold + 1\n    \n    trainer = pl.Trainer(\n#                         gpus=-1 if torch.cuda.is_available() else None, \n                        tpu_cores=[tpu_core],\n#                         precision=16 if torch.cuda.is_available() else 32,\n                        precision=16,\n#                         plugins='ddp_sharded',\n                        max_epochs=EPOCHS,\n                        checkpoint_callback=checkpoint_callback,\n                        callbacks=callbacks)\n    model = Model()\n    pl_dm = CassavaDataModule(fold=fold)\n    pl_module = CassavaPLModule(hparams={'lr':LR, 'batch_size':BATCH_SIZE}, model=model)\n    \n    trainer.use_native_amp = False\n    trainer.fit(pl_module, pl_dm)\n    \n    print(checkpoint_callback.best_model_path, checkpoint_callback.best_model_score)\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"import joblib as jl\nparallel = jl.Parallel(n_jobs=FOLDS, backend='threading', batch_size=1)\nparallel(jl.delayed(train)(i) for i in range(FOLDS))","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}