{"cells":[{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"from IPython.display import HTML\nfile = open(\"../input/notebookassets/custom.css\")\nHTML(\"<style>\"+file.read()+\"</style>\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Torch Multi-Model Trainer + Automatic Mixed Precision + Augmentations\n\nThis is the training pipeline that I've been using for a long time. Feel free to fork it, and make your changes!"},{"metadata":{},"cell_type":"markdown","source":"<p>If you like my work, please consider giving an <strong style=\"font-family:Verdana;color:red;font-size:25px\">Upvote</strong></p>"},{"metadata":{"execution":{"iopub.execute_input":"2021-01-19T08:27:59.100629Z","iopub.status.busy":"2021-01-19T08:27:59.100122Z","iopub.status.idle":"2021-01-19T08:27:59.103902Z","shell.execute_reply":"2021-01-19T08:27:59.103511Z"},"papermill":{"duration":0.020823,"end_time":"2021-01-19T08:27:59.104022","exception":false,"start_time":"2021-01-19T08:27:59.083199","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"import sys\nsys.path.append(\"../input/timm-pytorch-image-models/pytorch-image-models-master\")","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2021-01-19T08:27:59.213331Z","iopub.status.busy":"2021-01-19T08:27:59.212782Z","iopub.status.idle":"2021-01-19T08:28:02.925775Z","shell.execute_reply":"2021-01-19T08:28:02.924696Z"},"papermill":{"duration":3.735036,"end_time":"2021-01-19T08:28:02.925888","exception":false,"start_time":"2021-01-19T08:27:59.190852","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport warnings\n\nimport cv2\nfrom tqdm.notebook import tqdm\n\nfrom albumentations import (\n    HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,\n    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,\n    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine, RandomResizedCrop,\n    IAASharpen, IAAEmboss, RandomBrightnessContrast, Flip, OneOf, Compose, Normalize, Cutout, CoarseDropout, ShiftScaleRotate, CenterCrop, Resize\n)\n\nfrom albumentations.pytorch import ToTensorV2\n\nimport timm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.cuda.amp import autocast, GradScaler\nfrom torch.nn.modules.loss import _WeightedLoss\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau, CosineAnnealingWarmRestarts\n\nfrom sklearn.metrics import accuracy_score\nfrom sklearn.model_selection import StratifiedKFold\n\nwarnings.simplefilter(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","execution":{"iopub.execute_input":"2021-01-19T08:28:02.958214Z","iopub.status.busy":"2021-01-19T08:28:02.957544Z","iopub.status.idle":"2021-01-19T08:28:02.961648Z","shell.execute_reply":"2021-01-19T08:28:02.961241Z"},"papermill":{"duration":0.021514,"end_time":"2021-01-19T08:28:02.961739","exception":false,"start_time":"2021-01-19T08:28:02.940225","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"class Config:\n    CFG = {\n        'img_size': 384,\n        'tta': 3,\n        'wd': 1e-6\n    }","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2021-01-19T08:28:03.004476Z","iopub.status.busy":"2021-01-19T08:28:03.001682Z","iopub.status.idle":"2021-01-19T08:28:03.007154Z","shell.execute_reply":"2021-01-19T08:28:03.006635Z"},"papermill":{"duration":0.031928,"end_time":"2021-01-19T08:28:03.007239","exception":false,"start_time":"2021-01-19T08:28:02.975311","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"def plot_results(train_acc, valid_acc, train_loss, valid_loss, nb_epochs):\n    epochs = [i for i in range(nb_epochs)]\n    \n    fig, ax = plt.subplots(1, 2)\n    fig.set_size_inches(20, 10)\n    \n    ax[0].plot(epochs, train_acc, 'go-', label='Training Accuracy')\n    ax[0].plot(epochs, valid_acc, 'ro-', label='Validation Accuracy')\n    ax[0].set_title('Training & Validation Accuracy')\n    ax[0].legend()\n    ax[0].set_xlabel('Epochs')\n    ax[0].set_ylabel('Accuracy')\n    \n    ax[1].plot(epochs, train_loss, 'go-', label='Training Loss')\n    ax[1].plot(epochs, valid_loss, 'ro-', label='Validation Loss')\n    ax[1].set_title('Training & Validation Loss')\n    ax[1].legend()\n    ax[1].set_xlabel('Epochs')\n    ax[1].set_ylabel('Loss')\n    \n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Augmentations\nI've used some basic Augmentations, will probably be adding CutMix and Mixup in the future."},{"metadata":{"execution":{"iopub.execute_input":"2021-01-19T08:28:03.045396Z","iopub.status.busy":"2021-01-19T08:28:03.044724Z","iopub.status.idle":"2021-01-19T08:28:03.048403Z","shell.execute_reply":"2021-01-19T08:28:03.047983Z"},"papermill":{"duration":0.027367,"end_time":"2021-01-19T08:28:03.048481","exception":false,"start_time":"2021-01-19T08:28:03.021114","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"class Augments:\n    \"\"\"\n    Contains Train, Validation and Testing Augments\n    \"\"\"\n    train_augments = Compose([\n            RandomResizedCrop(Config.CFG['img_size'], Config.CFG['img_size']),\n            Transpose(p=0.5),\n            HorizontalFlip(p=0.5),\n            VerticalFlip(p=0.5),\n            ShiftScaleRotate(p=0.5),\n            HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),\n            RandomBrightnessContrast(brightness_limit=(-0.1,0.1), contrast_limit=(-0.1, 0.1), p=0.5),\n            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),\n            CoarseDropout(p=0.5),\n            Cutout(p=0.5),\n            ToTensorV2(p=1.0),\n        ],p=1.)\n    \n    valid_augments = Compose([\n            CenterCrop(Config.CFG['img_size'], Config.CFG['img_size'], p=1.),\n            Resize(Config.CFG['img_size'], Config.CFG['img_size']),\n            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0, p=1.0),\n            ToTensorV2(p=1.0),\n        ], p=1.)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model Classes\nRather than having all models inside one huge model class and then differentiating using `if-else` statements, I have made different classes for different families of models."},{"metadata":{"execution":{"iopub.execute_input":"2021-01-19T08:28:03.092315Z","iopub.status.busy":"2021-01-19T08:28:03.091153Z","iopub.status.idle":"2021-01-19T08:28:03.093707Z","shell.execute_reply":"2021-01-19T08:28:03.094135Z"},"papermill":{"duration":0.031825,"end_time":"2021-01-19T08:28:03.094232","exception":false,"start_time":"2021-01-19T08:28:03.062407","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"class ResNextModel(nn.Module):\n    \"\"\"\n    Model Class for ResNext Model Architectures\n    \"\"\"\n    def __init__(self, num_classes=5, model_name='resnext50d_32x4d', pretrained=True):\n        super(ResNextModel, self).__init__()\n        self.model = timm.create_model(model_name, pretrained=pretrained)\n        self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)\n        \n    def forward(self, x):\n        x = self.model(x)\n        return x\n\nclass ResNetModel(nn.Module):\n    \"\"\"\n    Model Class for ResNet Models\n    \"\"\"\n    def __init__(self, num_classes=5, model_name='resnet18', pretrained=True):\n        super(ResNetModel, self).__init__()\n        self.model = timm.create_model(model_name, pretrained=pretrained)\n        self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)\n        \n    def forward(self, x):\n        x = self.model(x)\n        return x\n\nclass EfficientNetModel(nn.Module):\n    \"\"\"\n    Model Class for EfficientNet Model\n    \"\"\"\n    def __init__(self, num_classes=5, model_name='efficientnet_b5', pretrained=True):\n        super(EfficientNetModel, self).__init__()\n        self.model = timm.create_model(model_name, pretrained=pretrained)\n        self.model.classifier = nn.Linear(self.model.classifier.in_features, num_classes)\n        \n    def forward(self, x):\n        x = self.model(x)\n        return x\n    \nclass VITModel(nn.Module):\n    \"\"\"\n    Model Class for VIT Model\n    \"\"\"\n    def __init__(self, num_classes=5, model_name='vit_base_patch16_224', pretrained=True):\n        super(VITModel, self).__init__()\n        self.model = timm.create_model(model_name, pretrained)\n        self.model.head = nn.Linear(self.model.head.in_features, num_classes)\n    \n    def forward(self, x):\n        x = self.model(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset Class\nJust a simple dataset class. I am taking the csv file and appending the right path on it. After that it's just reading and returning images after getting augmented along with corresponding labels."},{"metadata":{"execution":{"iopub.execute_input":"2021-01-19T08:28:03.134074Z","iopub.status.busy":"2021-01-19T08:28:03.132379Z","iopub.status.idle":"2021-01-19T08:28:03.13474Z","shell.execute_reply":"2021-01-19T08:28:03.135155Z"},"papermill":{"duration":0.027008,"end_time":"2021-01-19T08:28:03.135249","exception":false,"start_time":"2021-01-19T08:28:03.108241","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"class LDCData(Dataset):\n    def __init__(self, df, num_classes=5, is_train=True, augments=None, img_size=Config.CFG['img_size'], img_path=\"../input/cassava-leaf-disease-classification/train_images/\"):\n        super().__init__()\n        self.df = df.sample(frac=1).reset_index(drop=True)\n        self.num_classes = num_classes\n        self.is_train = is_train\n        self.augments = augments\n        self.img_size = img_size\n        self.img_path = img_path\n        \n        # Add the Right Image Path\n        self.df['image_id'] = self.df['image_id'].apply(lambda x: os.path.join(self.img_path, x))\n    \n    def __getitem__(self, idx):\n        # Read the image, Resize, convert to RGB from BGR\n        img = cv2.imread(self.df['image_id'][idx])\n        img = img[:, :, ::-1]\n        \n        # Augments must be albumentations\n        if self.augments:\n            img = self.augments(image=img)['image']\n        \n        if self.is_train:\n            label = self.df['label'][idx]\n            return img, label\n        \n        return img\n    \n    def __len__(self):\n        return len(self.df)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Trainer Class\nThis is the heart of this notebook. Rather than writing the whole training and validation code in one huge cell and then running it with dataloaders, I have went ahead and made a **very generalised** class that has intuitive functions for both training and validation. \n\nIt also displays both training and validation progress with the help of progress bars. Fork this notebook to get started!"},{"metadata":{"execution":{"iopub.execute_input":"2021-01-19T08:28:03.257347Z","iopub.status.busy":"2021-01-19T08:28:03.24086Z","iopub.status.idle":"2021-01-19T08:28:03.260035Z","shell.execute_reply":"2021-01-19T08:28:03.259586Z"},"papermill":{"duration":0.041171,"end_time":"2021-01-19T08:28:03.260114","exception":false,"start_time":"2021-01-19T08:28:03.218943","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"class Trainer:\n    def __init__(self, train_dataloader, valid_dataloader, model, optimizer, loss_fn, val_loss_fn, scheduler, device=\"cuda:0\", plot_results=True):\n        \"\"\"\n        TODO: Implement the ROC-AUC Scheduler stuff\n        \"\"\"\n        self.train = train_dataloader\n        self.valid = valid_dataloader\n        self.optim = optim\n        self.loss_fn = loss_fn\n        self.val_loss_fn = val_loss_fn\n        self.scheduler = scheduler\n        self.device = device\n        self.plot_results = plot_results\n    \n    def train_one_cycle(self):\n        \"\"\"\n        Runs one epoch of training, backpropagation, optimization and gets train accuracy\n        \"\"\"\n        model.train()\n        train_prog_bar = tqdm(self.train, total=len(self.train))\n\n        all_train_labels = []\n        all_train_preds = []\n        \n        running_loss = 0\n        \n        for xtrain, ytrain in train_prog_bar:\n            xtrain = xtrain.to(device).float()\n            ytrain = ytrain.to(device).long()\n            \n            with autocast():\n                # Get predictions\n                z = model(xtrain)\n\n                # Training\n                train_loss = self.loss_fn(z, ytrain)\n                scaler.scale(train_loss).backward()\n                \n                scaler.step(self.optim)\n                scaler.update()\n                self.optim.zero_grad()\n\n                # For averaging and reporting later\n                running_loss += train_loss\n\n                # Convert the predictions and corresponding labels to right form\n                train_predictions = torch.argmax(z, 1).detach().cpu().numpy()\n                train_labels = ytrain.detach().cpu().numpy()\n\n                # Append current predictions and current labels to a list\n                all_train_labels += [train_predictions]\n                all_train_preds += [train_labels]\n\n            # Show the current loss to the progress bar\n            train_pbar_desc = f'loss: {train_loss.item():.4f}'\n            train_prog_bar.set_description(desc=train_pbar_desc)\n        \n        # After all the batches are done, calculate the training accuracy\n        all_train_preds = np.concatenate(all_train_preds)\n        all_train_labels = np.concatenate(all_train_labels)\n        \n        train_acc = (all_train_preds == all_train_labels).mean()\n        print(f\"Training Accuracy: {train_acc:.4f}\")\n        \n        # Now average the running loss over all batches and return\n        train_running_loss = running_loss / len(self.train)\n        \n        # Free up memory\n        del all_train_labels, all_train_preds, train_predictions, train_labels, xtrain, ytrain, z\n        \n        return (train_acc, train_running_loss)\n\n    def valid_one_cycle(self):\n        \"\"\"\n        Runs one epoch of prediction and validation accuracy calculation\n        \"\"\"        \n        model.eval()\n        \n        valid_prog_bar = tqdm(self.valid, total=len(self.valid))\n        \n        with torch.no_grad():\n            all_valid_labels = []\n            all_valid_preds = []\n            \n            running_loss = 0\n            \n            for xval, yval in valid_prog_bar:\n                xval = xval.to(device).float()\n                yval = yval.to(device).long()\n                \n                val_z = model(xval)\n                \n                val_loss = self.val_loss_fn(val_z, yval)\n                \n                running_loss += val_loss.item()\n                \n                val_pred = torch.argmax(val_z, 1).detach().cpu().numpy()\n                val_label = yval.detach().cpu().numpy()\n                \n                all_valid_labels += [val_label]\n                all_valid_preds += [val_pred]\n            \n                # Show the current loss\n                valid_pbar_desc = f\"loss: {val_loss.item():.4f}\"\n                valid_prog_bar.set_description(desc=valid_pbar_desc)\n            \n            # Get the final loss\n            final_loss_val = running_loss / len(self.valid)\n            \n            # Get Validation Accuracy\n            all_valid_labels = np.concatenate(all_valid_labels)\n            all_valid_preds = np.concatenate(all_valid_preds)\n            \n            val_accuracy = (all_valid_preds == all_valid_labels).mean()\n            print(f\"Validation Accuracy: {val_accuracy:.4f}\")\n            \n            # Free up memory\n            del all_valid_labels, all_valid_preds, val_label, val_pred, xval, yval, val_z\n            \n        return (val_accuracy, final_loss_val, model)","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2021-01-19T08:28:03.323469Z","iopub.status.busy":"2021-01-19T08:28:03.322223Z","iopub.status.idle":"2021-01-19T08:28:03.325082Z","shell.execute_reply":"2021-01-19T08:28:03.324645Z"},"papermill":{"duration":0.021208,"end_time":"2021-01-19T08:28:03.325168","exception":false,"start_time":"2021-01-19T08:28:03.30396","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"nb_epochs = 5\ndevice = torch.device(\"cuda\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Running everything together\nNow, this is the main running code. \n\nI have split the csv files of the dataset already in 5 folds and uploaded it as a seperate dataset. This way you can use it too!\nHere I am just initializing model, optimizers, losses and trainer instance.\nLater, I just call the training instance and then traing the model.\n\nAt the end, I am just plotting the results."},{"metadata":{"execution":{"iopub.execute_input":"2021-01-19T08:28:03.375672Z","iopub.status.busy":"2021-01-19T08:28:03.375142Z","iopub.status.idle":"2021-01-19T11:43:27.63887Z","shell.execute_reply":"2021-01-19T11:43:27.639317Z"},"papermill":{"duration":11724.299182,"end_time":"2021-01-19T11:43:27.639494","exception":false,"start_time":"2021-01-19T08:28:03.340312","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"# Let's try individual fold modelling\nfold = 0\ntrain_fold = pd.read_csv(f\"../input/cassava-5folds-data/fold_{fold}_train.csv\")\nvalid_fold = pd.read_csv(f\"../input/cassava-5folds-data/fold_{fold}_valid.csv\")\n\ntrain_set = LDCData(df=train_fold, augments=Augments.train_augments)\nvalid_set = LDCData(df=valid_fold, augments=Augments.valid_augments)\n\ntrain = DataLoader(\n    train_set,\n    batch_size=16,\n    shuffle=True,\n    pin_memory=False,\n    drop_last=False,\n    num_workers=8\n)\n\nvalid = DataLoader(\n    valid_set,\n    batch_size=32,\n    shuffle=False,\n    pin_memory=False,\n    num_workers=8\n)\n\nmodel = VITModel(num_classes=5, model_name='vit_base_patch16_384').to(device)\noptim = torch.optim.AdamW(model.parameters(), lr=1e-5, weight_decay=Config.CFG['wd'])\nloss_fn = nn.CrossEntropyLoss().to(device)\nloss_fn_val = nn.CrossEntropyLoss().to(device)\n\ntrainer = Trainer(\n    train_dataloader=train,\n    valid_dataloader=valid,\n    model=model,\n    optimizer=optim,\n    loss_fn=loss_fn,\n    val_loss_fn=loss_fn_val,\n    scheduler=None,\n    device=device,\n)\n\ntrain_accs = []\nvalid_accs = []\ntrain_losses = []\nvalid_losses = []\n\nscaler = GradScaler()\n\nfor epoch in range(nb_epochs):\n    print(f\"{'-'*20} EPOCH: {epoch}/{nb_epochs} {'-'*20}\")\n\n    # Run one training epoch\n    current_train_acc, current_train_loss = trainer.train_one_cycle()\n    train_accs.append(current_train_acc)\n    train_losses.append(current_train_loss)\n\n    # Run one validation epoch\n    current_val_acc, current_val_loss, op_model = trainer.valid_one_cycle()\n    valid_accs.append(current_val_acc)\n    valid_losses.append(current_val_loss)\n\n    # Empty CUDA cache\n    torch.cuda.empty_cache()\n    \n    # Save the model every epoch\n    print(f\"Saving Model for this epoch...\")\n    torch.save(op_model.state_dict(), f\"vit_base_p16_384_fold_{fold}_model.pth\")\n    \ndel train_set, valid_set, train, valid, model, optim, loss_fn, loss_fn_val, trainer, scaler\ntorch.cuda.empty_cache()\n\nplot_results(train_accs, valid_accs, train_losses, valid_losses, nb_epochs)","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}