{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":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 1.7 --apt-packages libomp5 libopenblas-dev\n!pip install pytorch-lightning --quiet","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import torch\nimport torchvision\nimport pytorch_lightning as pl\nfrom torch.utils.data import DataLoader, Dataset\nfrom torchvision.utils import make_grid\nimport torchvision.transforms as T\nfrom torch import nn, optim\nimport torch.nn.functional as F\n\nimport glob\nimport os\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport pandas as pd\nimport warnings\nimport cv2\nfrom PIL import Image\nfrom pathlib import Path\n\nimport io\nimport tensorflow as tf\nfrom tqdm.auto import tqdm\nimport logging\n\n\nwarnings.filterwarnings(\"ignore\")\npd.set_option(\"display.max_colwidth\", None)\nlogger = logging.getLogger(\"lightning\")\npl.seed_everything(123)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN_PATH = Path(\"../input/tpu-getting-started/tfrecords-jpeg-512x512/train/\")\nVALID_PATH = Path(\"../input/tpu-getting-started/tfrecords-jpeg-512x512/val/\")\nTEST_PATH  = Path(\"../input/tpu-getting-started/tfrecords-jpeg-512x512/test/\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"CLASSES = ['pink primrose',    'hard-leaved pocket orchid', 'canterbury bells', 'sweet pea',     'wild geranium',     'tiger lily',           'moon orchid',              'bird of paradise', 'monkshood',        'globe thistle',      \n           'snapdragon',       \"colt's foot\",               'king protea',      'spear thistle', 'yellow iris',       'globe-flower',         'purple coneflower',        'peruvian lily',    'balloon flower',   'giant white arum lily', \n           'fire lily',        'pincushion flower',         'fritillary',       'red ginger',    'grape hyacinth',    'corn poppy',           'prince of wales feathers', 'stemless gentian', 'artichoke',        'sweet william',         \n           'carnation',        'garden phlox',              'love in the mist', 'cosmos',        'alpine sea holly',  'ruby-lipped cattleya', 'cape flower',              'great masterwort', 'siam tulip',       'lenten rose',           \n           'barberton daisy',  'daffodil',                  'sword lily',       'poinsettia',    'bolero deep blue',  'wallflower',           'marigold',                 'buttercup',        'daisy',            'common dandelion',      \n           'petunia',          'wild pansy',                'primula',          'sunflower',     'lilac hibiscus',    'bishop of llandaff',   'gaura',                    'geranium',         'orange dahlia',    'pink-yellow dahlia',    \n           'cautleya spicata', 'japanese anemone',          'black-eyed susan', 'silverbush',    'californian poppy', 'osteospermum',         'spring crocus',            'iris',             'windflower',       'tree poppy',            \n           'gazania',          'azalea',                    'water lily',       'rose',          'thorn apple',       'morning glory',        'passion flower',           'lotus',            'toad lily',        'anthurium',             \n           'frangipani',       'clematis',                  'hibiscus',         'columbine',     'desert-rose',       'tree mallow',          'magnolia',                 'cyclamen ',        'watercress',       'canna lily',            \n           'hippeastrum ',     'bee balm',                  'pink quill',       'foxglove',      'bougainvillea',     'camellia',             'mallow',                   'mexican petunia',  'bromelia',         'blanket flower', \n           'trumpet creeper',  'blackberry lily',           'common tulip',     'wild rose']       ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_feature = {\n    \"class\": tf.io.FixedLenFeature([], tf.int64),\n    \"id\"   : tf.io.FixedLenFeature([], tf.string),\n    \"image\": tf.io.FixedLenFeature([], tf.string),\n}\n\ntest_feature = {\n    'id'   : tf.io.FixedLenFeature([], tf.string),\n    'image': tf.io.FixedLenFeature([], tf.string),\n}\n\n\ndef parse_train_image(example_proto):\n    \"\"\"parse a single train tfrecord example\"\"\"\n    return tf.io.parse_single_example(example_proto, train_feature)\n\ndef parse_test_image(example_proto):\n    \"\"\"parse a single test tfrecord example\"\"\"\n    return tf.io.parse_single_example(example_proto, test_feature)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_files  = glob.glob(str(TRAIN_PATH/\"*.tfrec\"))\ntrain_class  = []\ntrain_images = []\ntrain_ids    = []\n\nprint(f\"Reading tfrecords from {TRAIN_PATH}\")\nfor i in tqdm(train_files):\n    train_image_dataset = tf.data.TFRecordDataset(i)\n    \n    train_image_dataset = train_image_dataset.map(parse_train_image)\n\n    ids = [str(id_features['id'].numpy())[2:-1] for id_features in train_image_dataset]\n    train_ids = train_ids + ids\n\n    classes = [int(class_features['class'].numpy()) for class_features in train_image_dataset]\n    train_class = train_class + classes\n\n    images = [image_features['image'].numpy() for image_features in train_image_dataset]\n    train_images = train_images + images","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"val_files  = glob.glob(str(VALID_PATH/\"*.tfrec\"))\nval_ids    = []\nval_class  = []\nval_images = []\n\nprint(f\"Reading tfrecords from {VALID_PATH}\")\nfor i in tqdm(val_files):\n    val_image_dataset = tf.data.TFRecordDataset(i)\n    \n    val_image_dataset = val_image_dataset.map(parse_train_image)\n\n    ids = [str(id_features['id'].numpy())[2:-1] for id_features in val_image_dataset]\n    val_ids = val_ids + ids\n\n    classes = [int(class_features['class'].numpy()) for class_features in val_image_dataset]\n    val_class = val_class + classes\n\n    images = [image_features['image'].numpy() for image_features in val_image_dataset]\n    val_images = val_images + images","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_files  = glob.glob(str(TEST_PATH/\"*.tfrec\"))\ntest_ids    = []\ntest_images = []\n\nprint(f\"Reading tfrecords from {VALID_PATH}\")\nfor i in tqdm(test_files):\n    test_image_dataset = tf.data.TFRecordDataset(i)\n    test_image_dataset = test_image_dataset.map(parse_test_image)\n    ids = [str(id_features['id'].numpy())[2:-1] for id_features in test_image_dataset]\n    test_ids = test_ids + ids\n    images = [image_features['image'].numpy() for image_features in test_image_dataset]\n    test_images = test_images + images","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"assert len(train_ids) == len(train_class) == len(train_images)\nassert len(val_ids)   == len(val_class)   == len(val_images)\nassert len(test_ids)  == len(test_images)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class FlowersDataset(Dataset):\n    def __init__(self, ids, images, augments, classes = None, is_test = False):\n        self.ids     = ids\n        self.images  = images\n        self.classes = classes\n        self.augs    = augments\n        self.is_test = is_test\n        \n    def __len__(self):\n        return len(self.ids)\n    \n    def __getitem__(self, index):\n        idx  = self.ids[index]\n        image = Image.open(io.BytesIO(self.images[index]))\n        image = self.augs(image)\n        \n        if not self.is_test :\n            clas = torch.tensor(int(self.classes[index]))\n            return idx, image, clas \n        \n        elif self.is_test:\n            return idx, image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class FlowersDataModule(pl.LightningDataModule):\n    def __init__(self, train_vars:tuple, test_vars:tuple, valid_vars:tuple, batch_size:int = 256, input_dims:int = 512):\n        super().__init__()\n        self.train_ids, self.train_images, self.train_classes  = train_vars\n        self.valid_ids, self.valid_images, self.valid_classes  = valid_vars\n        self.test_ids,  self.test_images = test_vars\n        \n        # Imagenet means and stds\n        mean = [0.485, 0.456, 0.406]\n        std  = [0.229, 0.224, 0.225]\n        \n        self.batch_size = batch_size\n        \n        self.train_augs = T.Compose([\n            T.Resize(size=(input_dims, input_dims)),\n            T.RandomHorizontalFlip(),\n            T.ToTensor(),\n            T.Normalize(mean=mean, std=std),\n        ])\n        \n          \n        self.valid_augs = T.Compose([\n            T.Resize(size=(input_dims, input_dims)),\n            T.ToTensor(),\n            T.Normalize(mean=mean, std=std),\n        ])\n        \n        self.test_augs = T.Compose([\n            T.Resize(size=(input_dims, input_dims)),\n            T.ToTensor(),\n            T.Normalize(mean=mean, std=std),\n        ])\n    \n    def setup(self, stage=None):\n        if stage == 'fit' or stage is None:\n            self.flowers_train = FlowersDataset(self.train_ids, self.train_images, self.train_augs, self.train_classes)\n            self.flowers_valid = FlowersDataset(self.valid_ids, self.valid_images, self.valid_augs, self.valid_classes)\n        \n        if stage == 'test' or stage is None:\n            self.flowers_test  = FlowersDataset(self.test_ids,  self.test_images, augments=self.test_augs, is_test = True)\n        \n    def train_dataloader(self):\n        return DataLoader(self.flowers_train, shuffle=True, batch_size=self.batch_size, pin_memory=True)\n    \n    def val_dataloader(self):\n        return DataLoader(self.flowers_valid, shuffle=False, batch_size=self.batch_size, pin_memory=True)\n    \n    def test_dataloader(self):\n        return DataLoader(self.flowers_test,  shuffle=False, batch_size=self.batch_size, pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class FlowerClassifier(pl.LightningModule):\n    def __init__(self, output_dims: int, learning_rate:float, weight_decay:float):\n        super().__init__()\n        self.save_hyperparameters()\n        \n        self.classifier  = torchvision.models.resnet34(pretrained=True, progress=True)\n        base_output_dims = self.classifier.fc.out_features\n        \n        self.lin1   = nn.Sequential(nn.BatchNorm1d(base_output_dims),  nn.Dropout(0.2), nn.ReLU(inplace=True))\n        self.lin2   = nn.Sequential(nn.Linear(base_output_dims, 1024), nn.BatchNorm1d(1024), nn.Dropout(0.5), nn.ReLU())\n        self.lin3   = nn.Sequential(nn.Linear(1024, 512),  nn.BatchNorm1d(512),  nn.Dropout(0.5), nn.ReLU())\n        self.output = nn.Sequential(nn.Linear(512, self.hparams.output_dims))\n        \n        self.accuracy = pl.metrics.Accuracy()\n\n        self.results    = pd.DataFrame()\n        self.test_idxs  = []\n        self.test_preds = []\n        \n    def forward(self, x):\n        out = self.classifier(x)\n        out = self.lin3(self.lin2(self.lin1(out)))\n        out = self.output(out)\n        return out\n    \n    def training_step(self, batch, batch_idx, *args, **kwargs):\n        _, image, clas = batch\n        y_hat = self(image)\n        loss  = F.cross_entropy(y_hat, clas)\n        self.log('loss', loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n        return loss\n        \n    def validation_step(self, batch, batch_idx, *args, **kwargs):\n        _, image, clas = batch\n        logits = self(image)\n        loss   = F.cross_entropy(logits, clas)\n        self.log(\"val_loss\", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n        \n        metric = self.accuracy(logits, clas)\n        self.log(\"accuracy\", metric, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n        \n    def test_step(self, batch, batch_idx, *args, **kwargs):\n        idx, image = batch\n        logits     = self(image)\n        # compute the output from the logits\n        _, preds   = logits.max(dim=1)\n        res        = list(preds.cpu().numpy())\n        \n        self.test_preds = self.test_preds + res\n        self.test_idxs  = self.test_idxs  + list(idx)\n        \n    def test_epoch_end(self, *args, **kwargs):\n        self.results[\"id\"]    = self.test_idxs\n        self.results[\"label\"] = self.test_preds\n        \n    def configure_optimizers(self, *args, **kwargs):\n        opt = optim.AdamW(self.parameters(), lr = self.hparams.learning_rate, weight_decay = self.hparams.weight_decay)\n        lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(opt, mode=\"min\")\n        \n        pl_scheduler = {\n            \"scheduler\": lr_scheduler, \n            \"interval\": \"epoch\", \n            \"frequency\": 1, \n            \"reduce_on_plateau\": True, \n            \"monitor\": \"val_loss\", \n            \"strict\": True\n        }\n        \n        return [opt], [pl_scheduler] ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Cross-Check DataModule\ntrain_vars = (train_ids, train_images, train_class)\nvalid_vars = (val_ids, val_images, val_class)\ntest_vars  = (test_ids, test_images)\n\nfake_dm = FlowersDataModule(train_vars, test_vars, valid_vars, batch_size=8, input_dims=224)\n\nfake_dm.setup(\"fit\")\ntrn_ds, val_ds = fake_dm.train_dataloader(), fake_dm.val_dataloader()\nfake_dm.setup(\"test\")\ntest_ds = fake_dm.test_dataloader()\n\ntrn_batch = next(iter(trn_ds))\n_, trn_image, _ = trn_batch\n\nval_batch = next(iter(val_ds))\n_, val_image, _ = val_batch\n\ntest_batch = next(iter(test_ds))\n_, test_image = test_batch\n\ntrn_grid  = make_grid(trn_image, normalize=True, nrow=4).permute(1, 2, 0).numpy()\nval_grid  = make_grid(val_image, normalize=True, nrow=4).permute(1, 2, 0).numpy()\ntest_grid = make_grid(test_image, normalize=True, nrow=4).permute(1, 2, 0).numpy()\n\n\nfig, (ax1, ax2, ax3) = plt.subplots(nrows=3, ncols=1, figsize=(15, 15))\n\nax1.axis(\"off\")\nax1.imshow(trn_grid);\n\nax2.axis(\"off\")\nax2.imshow(val_grid);\n\n\nax3.axis(\"off\")\nax3.imshow(test_grid);","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"logger = pl.loggers.CSVLogger(save_dir=\"/kaggle/working/\", name=\"kaggle_tpu_flowers\", version=\"001\")\n\ncallbacks = [\n    pl.callbacks.EarlyStopping(monitor=\"val_loss\", patience=5),\n    pl.callbacks.LearningRateMonitor(\"step\"),\n]\n\ntrainer = pl.Trainer(tpu_cores=8, precision=16, callbacks=callbacks,logger=logger,  gradient_clip_val=0.5,)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataModule = FlowersDataModule(train_vars, test_vars, valid_vars, batch_size = 8, input_dims = 224)\nmodel = FlowerClassifier(len(CLASSES), learning_rate = 0.002, weight_decay = 0.01)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trainer.fit(model, datamodule=dataModule)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trainer.test(model, datamodule=dataModule)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"classifier.results.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}