{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Welcome 🤓📚\n\nIn this Notebook, you will learn how to write a custom training loop in pure PyTorch, create custom torch `Dataset` class, compute metrics for model performance, and Scale the Training on any hardware like GPU, TPU, IPU or Distributed Training with [LightningLite](https://devblog.pytorchlightning.ai/scale-your-pytorch-code-with-lightninglite-d5692a303f00).","metadata":{}},{"cell_type":"markdown","source":"## 🕵 Explore the provided data\n\n(EDA is taken from [Notebook](https://www.kaggle.com/code/jirkaborovec/whale-dolphin-eda-classify-lit-flash?scriptVersionId=89274025) of Jirka)","metadata":{}},{"cell_type":"code","source":"!ls -l /kaggle/input/happy-whale-and-dolphin\n\nPATH_DATASET = \"/kaggle/input/happy-whale-and-dolphin\"","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-04-04T11:27:50.589618Z","iopub.execute_input":"2022-04-04T11:27:50.590216Z","iopub.status.idle":"2022-04-04T11:27:51.279592Z","shell.execute_reply.started":"2022-04-04T11:27:50.590131Z","shell.execute_reply":"2022-04-04T11:27:51.278896Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Browsing the metadata","metadata":{}},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport seaborn as sn\nimport matplotlib.pyplot as plt\n\nsn.set()\n\ndf_train = pd.read_csv(os.path.join(PATH_DATASET, \"train.csv\"))\ndisplay(df_train.head())\nprint(f\"Dataset size: {len(df_train)}\")\nprint(f\"Unique ids: {len(df_train['individual_id'].unique())}\")","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:27:51.281506Z","iopub.execute_input":"2022-04-04T11:27:51.281843Z","iopub.status.idle":"2022-04-04T11:27:52.210998Z","shell.execute_reply.started":"2022-04-04T11:27:51.281805Z","shell.execute_reply":"2022-04-04T11:27:52.209921Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Lets see how many speaced we have in the database...","metadata":{}},{"cell_type":"code","source":"counts_imgs = df_train[\"species\"].value_counts()\ncounts_inds = df_train.drop_duplicates(\"individual_id\")[\"species\"].value_counts()\n\nax = pd.concat({\"per Images\": counts_imgs, \"per Individuals\": counts_inds}, axis=1).plot.barh(grid=True, figsize=(7, 10))\nax.set_xscale('log')","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:27:52.212293Z","iopub.execute_input":"2022-04-04T11:27:52.212545Z","iopub.status.idle":"2022-04-04T11:27:53.584072Z","shell.execute_reply.started":"2022-04-04T11:27:52.212509Z","shell.execute_reply":"2022-04-04T11:27:53.583383Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"And compare they with unique individuals... \n\n**Note:** that the counts are in log scale","metadata":{}},{"cell_type":"code","source":"import numpy as np\nfrom pprint import pprint\n\nspecies_individuals = {}\nfor name, dfg in df_train.groupby(\"species\"):\n    species_individuals[name] = dfg[\"individual_id\"].value_counts()\n\nsi_max = max(list(map(len, species_individuals.values())))\nsi = {n: [0] * si_max for n in species_individuals}\nfor n, counts in species_individuals.items():\n    si[n][:len(counts)] = list(np.log(counts))\nsi = pd.DataFrame(si)","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:27:53.585362Z","iopub.execute_input":"2022-04-04T11:27:53.585716Z","iopub.status.idle":"2022-04-04T11:27:53.676340Z","shell.execute_reply.started":"2022-04-04T11:27:53.585675Z","shell.execute_reply":"2022-04-04T11:27:53.675668Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import seaborn as sn\n\nfig = plt.figure(figsize=(10, 8))\nax = sn.heatmap(si[:500].T, cmap=\"BuGn\", ax=fig.gca())","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:27:53.682403Z","iopub.execute_input":"2022-04-04T11:27:53.683081Z","iopub.status.idle":"2022-04-04T11:27:54.763518Z","shell.execute_reply.started":"2022-04-04T11:27:53.683041Z","shell.execute_reply":"2022-04-04T11:27:54.762698Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"And see the top individulas","metadata":{}},{"cell_type":"code","source":"ax = df_train[\"individual_id\"].value_counts(ascending=True)[-50:].plot.barh(figsize=(3, 8), grid=True)  # ascending=True","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:27:54.764760Z","iopub.execute_input":"2022-04-04T11:27:54.765179Z","iopub.status.idle":"2022-04-04T11:27:56.389783Z","shell.execute_reply.started":"2022-04-04T11:27:54.765139Z","shell.execute_reply":"2022-04-04T11:27:56.389101Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Browse some images","metadata":{}},{"cell_type":"code","source":"nb_species = len(df_train[\"species\"].unique())\nfig, axarr = plt.subplots(ncols=5, nrows=nb_species, figsize=(12, nb_species * 2))\n\nfor i, (name, dfg) in enumerate(df_train.groupby(\"species\")):\n    axarr[i, 0].set_title(name)\n    for j, (_, row) in enumerate(dfg[:5].iterrows()):\n        im_path = os.path.join(PATH_DATASET, \"train_images\", row[\"image\"])\n        img = plt.imread(im_path)\n        axarr[i, j].imshow(img)\n        axarr[i, j].set_axis_off()","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:27:56.393843Z","iopub.execute_input":"2022-04-04T11:27:56.395823Z","iopub.status.idle":"2022-04-04T11:29:09.914691Z","shell.execute_reply.started":"2022-04-04T11:27:56.395779Z","shell.execute_reply":"2022-04-04T11:29:09.913789Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install -q -U timm pytorch-lightning>=1.6","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2022-04-04T11:29:09.916023Z","iopub.execute_input":"2022-04-04T11:29:09.916477Z","iopub.status.idle":"2022-04-04T11:29:21.322020Z","shell.execute_reply.started":"2022-04-04T11:29:09.916432Z","shell.execute_reply":"2022-04-04T11:29:21.321171Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\n\nimport timm\nfrom PIL import Image\nfrom torchvision import transforms as T\nfrom torch.utils.data import Dataset, DataLoader, random_split\n\n\nfrom torchmetrics import F1\nimport torch.nn.functional as F\n\nfrom tqdm.auto import tqdm\n\nimport torch\nfrom torch.optim.lr_scheduler import StepLR","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:29:21.324428Z","iopub.execute_input":"2022-04-04T11:29:21.324707Z","iopub.status.idle":"2022-04-04T11:29:24.092171Z","shell.execute_reply.started":"2022-04-04T11:29:21.324667Z","shell.execute_reply":"2022-04-04T11:29:24.091278Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Create DataLoader 🗄️\n\n`DataLoader` is an iterable object which contains your input image data and the target label. To create a DataLoader, we first need to implement a torch `Dataset` class. We define `MyDataset` class which inherits from `Dataset` and it will implement `__len__` and `__getitem__` method.","metadata":{}},{"cell_type":"code","source":"label_to_idx = {e:i for i, e in enumerate(df_train.species.unique())}\n\nclass MyDataset(Dataset):\n    def __init__(self, df, transforms=None):\n        super().__init__()\n        self.df = df\n        self.root = os.path.join(PATH_DATASET, \"train_images\")\n        self.transforms = transforms\n    \n    def __len__(self): return len(self.df)\n    \n    def __getitem__(self, idx):\n        \n        data = self.df.iloc[idx]\n        image = Image.open(self.root + f\"/{data.image}\").convert(\"RGB\")\n        \n        if self.transforms:\n            image = self.transforms(image)\n        \n        label = label_to_idx[data.species]\n        return image, label\n\n\ndef random_split_dataset(data: Dataset, pct=0.9):\n    \"\"\"\n    Randomly splits dataset into two sets. Length of first split is len(data) * pct.\n    Source: https://github.com/gradsflow/gradsflow/blob/main/gradsflow/data/common.py#L20\n    Args:\n        data: pytorch Dataset object with `__len__` implementation.\n        pct: percentage of split.\n    \"\"\"\n    n = len(data)\n    split_1 = int(n * pct)\n    split_2 = n - split_1\n    return random_split(data, (split_1, split_2))","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:29:24.093968Z","iopub.execute_input":"2022-04-04T11:29:24.094231Z","iopub.status.idle":"2022-04-04T11:29:24.109818Z","shell.execute_reply.started":"2022-04-04T11:29:24.094190Z","shell.execute_reply":"2022-04-04T11:29:24.109174Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"We define image augmentation to make our classifier robust. We will use the function `random_split_dataset` to split dataset into train and validation set.\nOnce we have our Dataset object, we can create a DataLoader class like this - `dataloader = DataLoader(dataset, batch_size=8)`","metadata":{}},{"cell_type":"code","source":"transforms = T.Compose([T.AutoAugment(),\n                        T.Resize((224,224)), T.ToTensor()])\n\n\nds = MyDataset(df_train, transforms)\nsplit_pct = 0.9\n\ntrain_ds, val_ds = random_split_dataset(ds, 0.9)","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:29:24.111223Z","iopub.execute_input":"2022-04-04T11:29:24.111489Z","iopub.status.idle":"2022-04-04T11:29:24.140599Z","shell.execute_reply.started":"2022-04-04T11:29:24.111453Z","shell.execute_reply":"2022-04-04T11:29:24.139983Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Write Training Loop 🛠️\n\nNow, we have our dataloader ready we can create our classifier and write training loop.\nA training loop consists of model prediction, loss computation, backward propagation and model weight update by the optimizer.\n\n\nWe will start with a basic training loop then will use `LightningLite` to enable multiple hardware, precision and distributed training.","metadata":{}},{"cell_type":"markdown","source":"First we will create model, optimizer, loss function and metrics.","metadata":{}},{"cell_type":"code","source":"batch_size = 4\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\ndry_run = True\nnum_epochs = 5\n\n\nmodel = timm.create_model(\"efficientnet_b0\", pretrained=True, num_classes=len(label_to_idx))\nmodel = model.to(device)\noptimizer = torch.optim.AdamW(model.parameters(), 1e-4)\ncriterion = torch.nn.CrossEntropyLoss()\nmetric = F1().to(device)\n\ntrain_loader = DataLoader(train_ds, batch_size=batch_size)\nval_loader = DataLoader(val_ds, batch_size=batch_size)","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:29:24.141614Z","iopub.execute_input":"2022-04-04T11:29:24.141930Z","iopub.status.idle":"2022-04-04T11:29:30.729518Z","shell.execute_reply.started":"2022-04-04T11:29:24.141894Z","shell.execute_reply":"2022-04-04T11:29:30.728755Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training Loop\n\nFor writing the training loop we will iterate a `for loop` for given number of epochs `num_epochs`. Set the model to training mode with `model.train()` and iterate through the dataloader. We pass the data to model and calculate the crossentropy loss. We do `loss.backward()` to compute gradients followed by `optimizer.step()` which will update the model weights.\n\nFor model evaluation we define a validation loop which will calculate the `F1 accuracy` on the validation dataset. For validation we set our model to eval mode with `model.eval()` method. For calculating F1 accuracy, we use [TorchMetrics](https://torchmetrics.readthedocs.io/) which contains a collection of Machine Learning metrics for distributed, scalable PyTorch models and an easy-to-use API to create custom metrics.","metadata":{}},{"cell_type":"code","source":"# EPOCH LOOP\nfor epoch in tqdm(range(1, num_epochs + 1)):\n\n    # TRAINING LOOP\n    model.train()\n    for batch_idx, (data, target) in tqdm(enumerate(train_loader), total=len(train_ds)//batch_size):\n        data, target= data.to(device), target.to(device)\n        \n        optimizer.zero_grad()\n        output = model(data)\n        loss = criterion(output, target)\n        loss.backward()\n        optimizer.step()\n        \n        if (batch_idx == 0) or ((batch_idx + 1) % log_interval == 0):\n            print(\n                \"Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}\".format(\n                    epoch,\n                    batch_idx * len(data),\n                    len(train_loader.dataset),\n                    100.0 * batch_idx / len(train_loader),\n                    loss.item(),\n                )\n            )\n            if dry_run:\n                break\n\n\n    # TESTING LOOP\n    model.eval()\n    test_loss = 0\n    with torch.no_grad():\n        for data, target in val_loader:\n            data = data.to(device)\n            target = target.to(device)\n            output = model(data)\n            test_loss += criterion(output, target).item()\n\n            # WITH TorchMetrics\n            metric(output, target)\n\n            if dry_run:\n                break\n\n    # all_gather is used to aggregated the value across processes\n    test_loss = test_loss / len(val_loader.dataset)\n\n    print(f\"\\nTest set: Average loss: {test_loss:.4f}, Accuracy: ({metric.compute():.0f}%)\\n\")\n    metric.reset()\n\n    if dry_run:\n        break","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:29:30.730961Z","iopub.execute_input":"2022-04-04T11:29:30.731201Z","iopub.status.idle":"2022-04-04T11:29:38.383516Z","shell.execute_reply.started":"2022-04-04T11:29:30.731166Z","shell.execute_reply":"2022-04-04T11:29:38.382787Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# 👷 Scale Model Training \n\nOur dry run was successful 🎉!\nNow, let's scale our training on a hardware accelerator like GPU or TPU. We can also use distributed training if multiple devices are available. For this purpose we use `LightningLite`, it scales PyTorch model training loop with minimal changes. That means we will retain the full control of our training loop! It also enables Precision support abd DDP training.\n\nTo use LightningLite, we will import it from PyTorch Lightning library. We implement LightningLite and override `run` method. We can just copy paste our whole training loop code inside the `run` method and then just make these three changes to our code.\n\n1. `model, optimizer = self.setup(model, optimizer)`\n2. `dataloader = self.setup_dataloaders(dataloader)`\n3. Replace `loss.backward()` with `self.backward(loss)`\n\n\n![GIF](https://pl-public-data.s3.amazonaws.com/docs/static/images/lite/lightning_lite.gif)","metadata":{}},{"cell_type":"code","source":"from pytorch_lightning.lite import LightningLite\n\n# ref: github.com/pytorchlightning/pytorch-lightning/pl_examples\n\nclass CustomTrainer(LightningLite):\n        \n    def run(self, num_epochs, batch_size, gamma=0.7, dry_run: bool=False, save_model=True, log_interval=10):\n\n        model = timm.create_model(\"efficientnet_b0\", pretrained=True, num_classes=len(label_to_idx))\n        optimizer = torch.optim.AdamW(model.parameters(), 1e-4)\n        criterion = torch.nn.CrossEntropyLoss()\n        metric = F1().to(self.device)\n        print(self.device)\n        \n        # don't forget to call `setup` to prepare for model / optimizer for distributed training.\n        # the model is moved automatically to the right device.\n        model, optimizer = self.setup(model, optimizer)\n        \n        pin_memory = \"cuda\" in self.device.type\n        train_loader, val_loader = self.setup_dataloaders(DataLoader(train_ds, batch_size=batch_size, pin_memory=pin_memory),\n                                            DataLoader(val_ds, batch_size=batch_size, pin_memory=pin_memory))\n\n        scheduler = StepLR(optimizer, step_size=1, gamma=gamma)\n\n        # EPOCH LOOP\n        for epoch in tqdm(range(1, num_epochs + 1)):\n\n            # TRAINING LOOP\n            model.train()\n            for batch_idx, (data, target) in tqdm(enumerate(train_loader), total=len(train_ds)//batch_size):\n                # NOTE: no need to call `.to(device)` on the data, target\n                optimizer.zero_grad()\n                output = model(data)\n                loss = criterion(output, target)\n                self.backward(loss)  # instead of loss.backward()\n\n                optimizer.step()\n                if (batch_idx == 0) or ((batch_idx + 1) % log_interval == 0):\n                    print(\n                        \"Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}\".format(\n                            epoch,\n                            batch_idx * len(data),\n                            len(train_loader.dataset),\n                            100.0 * batch_idx / len(train_loader),\n                            loss.item(),\n                        )\n                    )\n                    if dry_run:\n                        break\n\n            scheduler.step()\n\n            # TESTING LOOP\n            model.eval()\n            test_loss = 0\n            with torch.no_grad():\n                for data, target in val_loader:\n                    # NOTE: no need to call `.to(device)` on the data, target\n                    output = model(data)\n                    test_loss += criterion(output, target).item()\n\n                    # WITH TorchMetrics\n                    metric(output, target)\n\n                    if dry_run:\n                        break\n\n            # all_gather is used to aggregated the value across processes\n            test_loss = self.all_gather(test_loss).sum() / len(val_loader.dataset)\n\n            print(f\"\\nTest set: Average loss: {test_loss:.4f}, Accuracy: ({metric.compute()})\\n\")\n            metric.reset()\n\n            if dry_run:\n                break\n\n        # When using distributed training, use `self.save`\n        # to ensure the current process is allowed to save a checkpoint\n        if save_model:\n            self.save(model, \"model.pt\")","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:29:38.386518Z","iopub.execute_input":"2022-04-04T11:29:38.387111Z","iopub.status.idle":"2022-04-04T11:29:38.404627Z","shell.execute_reply.started":"2022-04-04T11:29:38.387061Z","shell.execute_reply":"2022-04-04T11:29:38.403899Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"That's all we need to do. Now we can select any supported hardware, precision type, number of devices, or [training strategy](https://pytorch-lightning.readthedocs.io/en/latest/starter/lightning_lite.html#strategy).\n\nRun this cell to train the model on one GPU.","metadata":{}},{"cell_type":"code","source":"trainer = CustomTrainer(accelerator = \"gpu\", devices=1)\ntrainer.run(num_epochs=1, batch_size=128, dry_run=True)","metadata":{"execution":{"iopub.status.busy":"2022-04-04T11:30:55.049047Z","iopub.execute_input":"2022-04-04T11:30:55.049393Z","iopub.status.idle":"2022-04-04T11:31:31.602098Z","shell.execute_reply.started":"2022-04-04T11:30:55.049348Z","shell.execute_reply":"2022-04-04T11:31:31.601346Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Happy Training! ⚡️🎉\n\nCheckout the libraries used in this notebook - \n\n1. [TorchMetrics](https://torchmetrics.readthedocs.io/)\n2. [LightningLite](https://pytorch-lightning.readthedocs.io/en/latest/starter/lightning_lite.html) ","metadata":{"execution":{"iopub.status.busy":"2022-04-02T18:22:50.596271Z","iopub.execute_input":"2022-04-02T18:22:50.596594Z","iopub.status.idle":"2022-04-02T18:22:50.603682Z","shell.execute_reply.started":"2022-04-02T18:22:50.596563Z","shell.execute_reply":"2022-04-02T18:22:50.602096Z"}}}]}