{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Install & Import libraries"},{"metadata":{"trusted":true},"cell_type":"code","source":"# For Using TPU\n#!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 install pytorch-lightning","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install ../input/pikachu-libraries/timm-0.3.4-py3-none-any.whl\n!pip install ../input/pikachu-libraries/omegaconf-2.0.6-py3-none-any.whl","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport sys\nimport cv2\nimport numpy as np\nimport pandas as pd\nsys.path.append('../input/pikachu-libraries/pytorch-lightning-1.1.7/pytorch-lightning-1.1.7')\n\n\nimport pytorch_lightning as pl\nfrom pytorch_lightning.metrics import Accuracy\nfrom pytorch_lightning import loggers\nfrom pytorch_lightning import seed_everything\nfrom pytorch_lightning import Trainer\nfrom pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom sklearn.model_selection import StratifiedKFold\nimport albumentations as A\nimport timm\nfrom omegaconf import OmegaConf","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Config"},{"metadata":{"trusted":true},"cell_type":"code","source":"conf_dict = {'batch_size': 32, \n             'epoch': 1,\n             'model_name': 'tf_efficientnet_b0'}\nconf_base = OmegaConf.create(conf_dict)\nconf = conf_base\n\n#conf_cli = OmegaConf.from_cli()\n#conf = OmegaConf.merge(conf_base, conf_cli)\nprint(OmegaConf.to_yaml(conf))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class DogBreedDataset(Dataset):\n    def __init__(self, dataframe, data_dir, transform=None, test=False):\n        self.data = dataframe.reset_index(drop=True)\n        self.data_dir = data_dir\n        self.transform = transform\n        self.test = test\n        cv2.setNumThreads(0)\n\n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        img_name = os.path.join(self.data_dir, self.data.loc[idx, \"id\"] + \".\" + \"jpg\")\n        image = cv2.imread(img_name)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        image = self.transform(image=image)\n        image = torch.from_numpy(image[\"image\"].transpose(2, 0, 1))\n        if not self.test:\n            label = self.data.loc[idx, \"label\"]\n            return image, label\n        else:\n            return image","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data Module"},{"metadata":{"trusted":true},"cell_type":"code","source":"import pytorch_lightning as pl\nfrom torch.utils.data import DataLoader\nimport pandas as pd\nfrom sklearn.model_selection import StratifiedKFold\nimport albumentations as A\n\nclass DogBreedDataModule(pl.LightningDataModule):\n\n    def __init__(self, conf):\n        super().__init__()\n        self.conf = conf\n        \n\n    # OPTIONAL, called only on 1 GPU/machine(for download or tokenize)\n    def prepare_data(self):\n        pass\n\n    # OPTIONAL, called for every GPU/machine\n    def setup(self, stage=None):\n        if stage == 'fit':\n            df = pd.read_csv('../input/dog-breed-identification/labels.csv')\n            \n            df['label'] = df.groupby(['breed']).ngroup()\n            self.breed2id_dict = df[['breed','label']].sort_values(by='label').set_index('breed').to_dict()['label'] \n            self.id2breed_dict = df[['breed','label']].sort_values(by='label').set_index('label').to_dict()['breed']  \n\n            kf = StratifiedKFold(n_splits=5, shuffle=True, random_state=2021)\n            for fold, (train_index, val_index) in enumerate(kf.split(df.values, df[\"breed\"])):\n                df.loc[val_index, \"fold\"] = int(fold)\n            df[\"fold\"] = df[\"fold\"].astype(int)\n\n            train_df = df[df[\"fold\"] != 0]\n            valid_df = df[df[\"fold\"] == 0]\n\n            train_transform = A.Compose([\n                        A.RandomResizedCrop(height=256, width=256, scale=(0.8, 1.0), ratio=(1, 1), interpolation=1, always_apply=False, p=1.0),\n                        A.Flip(always_apply=False, p=0.5),\n                        A.RandomGridShuffle(grid=(4, 4), always_apply=False, p=1.0),\n                        A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=10, val_shift_limit=10, always_apply=False, p=0.5),\n                        A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, brightness_by_max=True, always_apply=False, p=0.5),\n                        A.GaussNoise(var_limit=(10.0, 50.0), mean=0, always_apply=False, p=0.5),\n                        A.Rotate(limit=90, interpolation=1, border_mode=4, value=None, mask_value=None, always_apply=False, p=0.5),\n                        A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.0, rotate_limit=45, interpolation=1, border_mode=4, value=255, mask_value=None, always_apply=False, p=0.5),\n                        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, always_apply=False, p=1.0),\n                        A.Cutout(num_holes=8, max_h_size=8, max_w_size=8, fill_value=0, always_apply=False, p=0.5),\n                        ])\n\n            valid_transform = A.Compose([\n                        A.Resize(height=256, width=256, interpolation=1, always_apply=False, p=1.0),\n                        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, always_apply=False, p=1.0)\n                        ])\n\n            self.train_dataset = DogBreedDataset(train_df, '../input/dog-breed-identification/train', transform=train_transform)\n            self.valid_dataset = DogBreedDataset(valid_df, '../input/dog-breed-identification/train', transform=valid_transform)\n            \n        elif stage == 'test':\n            test_df = pd.read_csv('../input/dog-breed-identification/sample_submission.csv')\n            test_transform = A.Compose([\n                        A.Resize(height=256, width=256, interpolation=1, always_apply=False, p=1.0),\n                        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, always_apply=False, p=1.0)\n                        ])            \n            self.test_dataset = DogBreedDataset(test_df, '../input/dog-breed-identification/test', transform=test_transform, test=True)\n        \n\n    def train_dataloader(self):\n        return DataLoader(self.train_dataset, batch_size=self.conf.batch_size, num_workers=4, shuffle=True, pin_memory=True, drop_last=True)\n\n    def val_dataloader(self):\n        return DataLoader(self.valid_dataset, batch_size=self.conf.batch_size, num_workers=4, shuffle=False, pin_memory=True, drop_last=True)\n\n    def test_dataloader(self):\n        return DataLoader(self.test_dataset, batch_size=self.conf.batch_size, num_workers=4, shuffle=False, pin_memory=True, drop_last=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Lightning Module"},{"metadata":{"trusted":true},"cell_type":"code","source":"class LitSystem(pl.LightningModule):\n    def __init__(self, conf):\n        super().__init__()\n        #self.conf = conf\n        self.save_hyperparameters(conf)\n        self.model = timm.create_model(model_name=self.hparams.model_name, num_classes=120, pretrained=False, in_chans=3)\n        self.criteria = torch.nn.CrossEntropyLoss()\n        self.accuracy = Accuracy()\n\n    def forward(self, x):\n        # use forward for inference/predictions\n        return self.model(x)\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)\n\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)\n        \n        return [optimizer], [scheduler]\n\n    def training_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self.model(x)\n        loss = self.criteria(y_hat, y)\n        \n        self.log('train_loss', loss, on_epoch=True)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self.model(x)\n        loss = self.criteria(y_hat, y)\n        \n        return {\n            \"val_loss\": loss,\n            \"y\": y,\n            \"y_hat\": y_hat\n            }\n    \n    def validation_epoch_end(self, outputs):\n        avg_val_loss = torch.stack([x[\"val_loss\"] for x in outputs]).mean()\n        y = torch.cat([x[\"y\"] for x in outputs]).cpu()\n        y_hat = torch.cat([x[\"y_hat\"] for x in outputs]).cpu()\n\n        #preds = np.argmax(y_hat, axis=1)\n\n        val_accuracy = self.accuracy(y_hat, y)\n\n        self.log('avg_val_loss', avg_val_loss)\n        self.log('val_acc', val_accuracy)\n\n    def test_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self.model(x)\n        loss = self.criteria(y_hat, y)\n        self.log('test_loss', loss)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train"},{"metadata":{"trusted":true},"cell_type":"code","source":"seed_everything(2021)\n\ntb_logger = loggers.TensorBoardLogger(save_dir='tb_log/')\ncsv_logger = loggers.CSVLogger(save_dir='csv_log/')\n\nlr_monitor = LearningRateMonitor(logging_interval='step')\ncheckpoint_callback = ModelCheckpoint(dirpath='ckpt', monitor='avg_val_loss', \n                                      save_last=True, save_top_k=5, mode='min', \n                                      save_weights_only=True, filename='{epoch}-{avg_val_loss:.2f}')\n\ndata_module = DogBreedDataModule(conf)\n\nlit_model = LitSystem(conf)\n\ntrainer = Trainer(\n    logger=[tb_logger, csv_logger],\n    callbacks=[lr_monitor, checkpoint_callback],\n    max_epochs=conf.epoch,\n    gpus=-1,\n    amp_backend='native',\n    amp_level='O2',\n    precision=16\n        )\n\ntrainer.fit(lit_model, data_module)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Inference for valid"},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_pytorch_model(ckpt_name, model, ignore_suffix='model'):\n    state_dict = torch.load(ckpt_name)[\"state_dict\"]\n    new_state_dict = {}\n    for k, v in state_dict.items():\n        name = k\n        if name.startswith(str(ignore_suffix)+\".\"):\n            name = name.replace(str(ignore_suffix)+\".\", \"\")  # remove `model.`\n        new_state_dict[name] = v\n    model.load_state_dict(new_state_dict)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_module = DogBreedDataModule(conf)\ndata_module.prepare_data()\ndata_module.setup('fit')\n\ntargets = []\npreds = []\n\n#lit_model.eval().cuda()\nmodel = timm.create_model(model_name=conf.model_name, num_classes=120, pretrained=False, in_chans=3)\nmodel = load_pytorch_model('ckpt/last.ckpt', model).eval().cuda()\n\nfor batch in data_module.val_dataloader():\n    x, y = batch\n    with torch.no_grad():\n        y_hat = model(x.cuda())\n    targets.append(y.detach().cpu().numpy())\n    preds.append(y_hat.detach().cpu().numpy().argmax(axis=1))\ntargets = np.concatenate(targets, axis=0)\npreds = np.concatenate(preds, axis=0)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Inference for test"},{"metadata":{"trusted":true},"cell_type":"code","source":"data_module.setup('test')\npreds = []\n\n#lit_model.eval().cuda()\nmodel = timm.create_model(model_name=conf.model_name, num_classes=120, pretrained=False, in_chans=3)\nmodel = load_pytorch_model('ckpt/last.ckpt', model).eval().cuda()\n\nfor batch in data_module.test_dataloader():\n    x = batch\n    with torch.no_grad():\n        y_hat = model(x.cuda())\n    preds.append(y_hat.detach().cpu().numpy().argmax(axis=1))\npreds = np.concatenate(preds, axis=0)\n","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}