{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Imports and GPU-Check\n\nFirst, let's import everything that we'll need:\n* PyTorch and its Torchvision library (for image-related data preparation)\n* Pandas for reading CSVs\n* Numpy for storing arrays\n* os for some path-related functions\n* PyTorch Lightning for the magic :)"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import torch\nimport torchvision\nimport torchvision.transforms as transforms\nimport torch.optim as optim\nimport time\nimport torch.nn.functional as F\nimport torch.nn as nn\nfrom torchvision import models\nimport pandas as pd\nimport numpy as np\nimport os\n\nimport pytorch_lightning as pl","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Here we're just going to check for CUDA so that we can run on the GPU - I **highly recommend** using a GPU for image-related ML tasks, as the CPU is not suited for this kind of training and therefore runs *wayyyy* too slowly for productive training."},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# The Briefest Data Exploration Ever\nLet's now take a quick look at our data - as we can see, each image has an id which corresponds to its **filename**, and then a label which corresponds to its **classification**."},{"metadata":{"trusted":true},"cell_type":"code","source":"data = pd.read_csv(\"/kaggle/input/cassava-leaf-disease-classification/train.csv\")\ndata.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data Preparation: Dataset and DataModule Creation\nHere's where we start getting to the good stuff.\n\nFirst, we'll build a custom CassavaDataset class using PyTorch's Dataset class, which requires that we define three functions: the initialization of a class instance (in the *init* function) as well as *len* and *getitem* functions. All we're doing in the *init* function is reading from the correct `.csv` file based on whether we need the training or test dataset, and then storing the image_ids and labels in arrays along with other information we might need (such as the image directory and the transforms we should apply to each image).\n\nThen, in the *getitem* stage we actually load the image, perform any transforms necessary, and then return it with the corresponding label. "},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\n\nclass CassavaDataset(Dataset):\n    \"\"\" Cassava Dataset \"\"\"\n    \n    def __init__(self, root_dir, transform=None, stage=None):\n        if (stage):\n            # We're in test stage then\n            csv_output = pd.read_csv(os.path.join(root_dir, \"sample_submission.csv\"))\n            self.images_dir = os.path.join(root_dir, \"test_images\")\n        else:\n            csv_output = pd.read_csv(os.path.join(root_dir, \"train.csv\"))\n            self.images_dir = os.path.join(root_dir, \"train_images\")\n        self.image_urls = np.asarray(csv_output[\"image_id\"])\n        self.labels = np.asarray(csv_output[\"label\"])\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.image_urls)\n    \n    def __getitem__(self, idx):\n        # Get and load image\n        image_path = os.path.join(self.images_dir, self.image_urls[idx])\n        image = Image.open(image_path)\n        # Perform transforms if any\n        if self.transform:\n            image = self.transform(image)\n        # Get label\n        label = self.labels[idx]\n        return image, label","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"This is where we start to use Pytorch Lightning. We'll build a DataModule here, which lets us easily construct the PyTorch DataLoaders we'll use during training. As you can see in the *setup* function, we use the previously defined CassavaDataset class to build our dataset, and then split it up into training, validation, and test sets which are then returned as DataLoaders through their respective functions. "},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data import random_split\nimport math\n\nclass CassavaDataModule(pl.LightningDataModule):\n    \"\"\" Cassava DataModule for Lightning \"\"\"\n    def __init__(self, root_dir, transform=None, batch_size=32):\n        super().__init__()\n        self.batch_size = batch_size\n        self.root_dir = root_dir\n        self.transform = transform\n        \n    def setup(self, stage=None):\n        cassava_full = CassavaDataset(self.root_dir, self.transform)\n        train_data_len = math.floor(len(cassava_full) * 0.7)\n        val_data_len = len(cassava_full) - train_data_len\n        # Create train and validation datasets\n        self.cassava_train, self.cassava_val = random_split(cassava_full, [train_data_len, val_data_len], generator=torch.Generator().manual_seed(42))\n        \n        # Create test dataset\n        self.cassava_test = CassavaDataset(self.root_dir, self.transform, stage=\"test\")\n        \n    def train_dataloader(self):\n        return DataLoader(self.cassava_train, batch_size=self.batch_size)\n    \n    def val_dataloader(self):\n        return DataLoader(self.cassava_val, batch_size=self.batch_size)\n    \n    def test_dataloader(self):\n        return DataLoader(self.cassava_test, batch_size=self.batch_size)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Finally, we just need to define the transforms we'll be using, and then create our DataModule!\n\nThe transforms are as follows:\n* Resize each image to 224x224 as that is the image size that our ResNet model (defined later) was trained on\n* Convert the image to a Tensor so that PyTorch can handle it\n* Normalize images according to either standard normalization or ImageNet Normalization\n\nBoth normalization techniques are meant to ensure that our pixels, across images, have a similar data distribution - this is meant to help training (read more [here](https://becominghuman.ai/image-data-pre-processing-for-neural-networks-498289068258)). While the standard normalization works for most cases, ImageNet normalization is sometimes applied to ensure that all incoming images have a similar data distribution to the images that ImageNet was trained on. Test both of these and see which one works best for you!"},{"metadata":{"trusted":true},"cell_type":"code","source":"transform = transforms.Compose([\n    transforms.Resize((224, 224)),\n    transforms.ToTensor(),\n    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # Standard Normalization\n#      transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) #ImageNet Normalization\n    ])\n\nroot_dir = \"/kaggle/input/cassava-leaf-disease-classification/\"\ncassava_data = CassavaDataModule(root_dir, transform, batch_size=64)\ncassava_data.setup()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model Initialization\nGreat - so we've processed our data, setup a DataModule and its DataLoaders, now what? It's time to construct our model class!\n\nFor a model, PyTorch Lightning requires us to define just 5 functions. First we have the initialization function, where we set a few variables that we'll use for training (the learning rate and the loss function), as well as construct the two important parts of our model: the ImageNet-trained ResNet50 model along with the additional Linear layer. This additional layer will be used for the actual classifying of the Cassava leaves, as it will take the image \"features\" extracted by ResNet and then make a classification after some additional processing.\n\nThe next few functions are used for different parts of training. First, *configure_optimizers* allows us to set the optimization function we'll be using for training - here we've chosen Adam, as it is often considered the best-performing optimizer on image tasks. Next, we define our *forward* function, which essentially just describes how we go from start to finish with our model - from the input image to a classification prediction output. You might be wondering - shouldn't there be a backprop function then too? Thankfully, PyTorch and PyTorch Lightning handle that for us, so we don't have to worry about that at all :)\n\nThen, the final two functions here describe the training process. In both, we receive a batch of images, process them with our *forward* function, and then calculate the loss between the predicted classes and the actual classes. In the *training_step* we return this calculated loss as PyTorch Lightning will automatically use that loss with our previously-configured optimizer to compute the backpropagation for the model; which, as mentioned earlier, Lightning handles for us."},{"metadata":{"trusted":true},"cell_type":"code","source":"import torchvision.models as models\n\nclass ImageNetModel(pl.LightningModule):\n    def __init__(self, learning_rate=1e-3):\n        super().__init__()\n        \n        # Set our learning rate\n        self.learning_rate = learning_rate\n        \n        num_target_classes = 5\n        self.feature_extractor = models.resnet50(pretrained=True)\n        self.feature_extractor.eval()\n        \n        # Use the pretrained model to classify cassava\n        self.classifier = nn.Linear(1000, num_target_classes)\n        \n        # Create loss function\n        self.loss_fn = torch.nn.CrossEntropyLoss()\n    \n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)\n        return optimizer\n    \n    def forward(self, input_data):\n        representations = self.feature_extractor(input_data)\n        preds = self.classifier(representations)\n        return preds\n    \n    def training_step(self, train_batch, batch_idx):\n        x, y = train_batch\n        predictions = self.forward(x)\n        loss = self.loss_fn(predictions, y)\n        return loss\n    \n    def validation_step(self, val_batch, batch_idx):\n        x, y = val_batch\n        predictions = self.forward(x)\n        loss = self.loss_fn(predictions, y)\n        self.log('val_loss', loss)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Initializing our Model & Trainer\nNow that we've defined our model class, we're done with setup! Yes, it really was that easy.\n\nFrom here, it's time to create an instance of the model and an instance of PyTorch's Trainer (which is where the magic happens). Note that we set the `auto_lr_find` parameter to `True`, which allows us to \"tune\" the learning rate to an appropriate level (which the Trainer will find for us!) and use that during training."},{"metadata":{"trusted":true},"cell_type":"code","source":"model = ImageNetModel()\n\ntrainer = pl.Trainer(gpus=1, auto_lr_find=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Tune the Learning Rate"},{"metadata":{"trusted":true},"cell_type":"code","source":"trainer.tune(model, cassava_data)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Fitting our Model to Our Data\nNow, we just fit the PyTorch LightningModule-based model to our previously-defined DataModule using the Trainer - and then we have our results! Here we only ran it for one epoch (which means that our model will see each image in the dataset only *once*), but you can definitely run it for more (and should!) to get better results.\n\nAdditionally, the Trainer saves training logs and a version of our model after each epoch - which allows us to easily inspect our results without having to mess around with the \"administrative\" side of ML. Isn't it great that we can just focus on the data and model without having to worry about tangential things like logging?"},{"metadata":{"trusted":true},"cell_type":"code","source":"trainer.fit(model, cassava_data)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Evaluating our Performance\nFinally, here's a simple way for us to check our model's performance at the end on the validation and test sets.\n\nConsidering we only ran our model for one epoch, I'd say this is pretty good!"},{"metadata":{"trusted":true},"cell_type":"code","source":"def evaluate_results(loader):\n    model.eval()\n    correct = 0\n    total = 0\n    \n    with torch.no_grad():\n        for x, y in iter(loader):        \n            x = x.to(device)\n            y = y.to(device)\n            preds = model(x)\n            _, predicted = torch.max(preds, 1)\n\n            correct += (predicted == y).sum().item()\n            total += len(y)\n    return (correct / total)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Check the Validation Results\nvalidation_loader = cassava_data.val_dataloader()\nevaluate_results(validation_loader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"test_loader = cassava_data.test_dataloader()\nevaluate_results(test_loader)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Finding Our Saved Model\nFinally, we can see that PyTorch Lightning did indeed save a checkpoint of our model after the first epoch"},{"metadata":{"trusted":true},"cell_type":"code","source":"for dirname, _, filenames in os.walk('/kaggle/working/lightning_logs/version_0/checkpoints'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))","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}