{"cells":[{"metadata":{},"cell_type":"markdown","source":"# **Simple PyTorch Training pipeline to train using SnapMix Augmentation.** \n\n*References :* \n* https://arxiv.org/abs/2012.04846\n* https://github.com/Shaoli-Huang/SnapMix\n\n\n* V15 : Base version\n* V17 : Added warmup epoch(1), reducing snapmix to 0.5, Moved back prop etc outside autocast(as suggested in docs)\n* V18 : Increased lr of backbone and running for 10 epochs now"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install timm","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"from __future__ import print_function, division\nimport random\nimport os\nimport torch\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nimport torchvision\nimport timm\nfrom torchvision import models as tvmodels\nfrom torch.cuda.amp import autocast, GradScaler\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, utils\nfrom tqdm import tqdm\nimport torch.nn.functional as F\nfrom sklearn.metrics import accuracy_score\n\nimport albumentations as A\nimport numpy as np\nimport cv2\nfrom sklearn.model_selection import GroupKFold, StratifiedKFold","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from albumentations import Compose\nfrom albumentations.pytorch import ToTensorV2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Ignore warnings\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_PATH = '../input/cassava-leaf-disease-classification/'\nNUM_FOLDS = 5\nbs = 32\n# Running only 5 epochs to test (Train more offline ^_^)\nEPOCHS = 10\nsz = 512\nSNAPMIX_ALPHA = 5.0\nSNAPMIX_PCT = 0.5\nGRAD_ACCUM_STEPS = 1\nTIMM_MODEL = 'resnet50'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# TODO : Play around with SNAPMIX_PCT, SNAPMIX_ALPHA and EPOCHS to converge for max accuracy","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True\n\nSEED = 1234\nseed_everything(SEED)    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Cassava Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaDataset(Dataset):\n    \"\"\"Cassava dataset.\"\"\"\n\n    def __init__(self, dataframe, root_dir, transforms=None):\n        \"\"\"\n        Args:\n            dataframe (string): dataframe train/valid\n            root_dir (string): Directory with all the images.\n            transform (callable, optional): Optional transform to be applied\n                on a sample.\n        \"\"\"\n        super().__init__()\n        self.dataframe = dataframe\n        self.root_dir = root_dir\n        self.transforms = transforms\n\n    def __len__(self):\n        return len(self.dataframe)\n    \n    def get_img_bgr_to_rgb(self, path):\n        im_bgr = cv2.imread(path)\n        im_rgb = im_bgr[:, :, ::-1]\n        return im_rgb\n\n    def __getitem__(self, idx):\n        if torch.is_tensor(idx):\n            idx = idx.tolist()\n        img_name = os.path.join(self.root_dir,\n                                self.dataframe.iloc[idx, 0])\n        image = self.get_img_bgr_to_rgb(img_name)\n        if self.transforms:\n            image = self.transforms(image=image)['image']\n        csv_row = self.dataframe.iloc[idx, 1:]\n        sample = {\n            'image': image, \n            'label': csv_row.label,\n        }\n        return sample","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(DATA_PATH + \"train.csv\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Transforms using albumentations"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_transforms():\n    return Compose([\n            A.RandomResizedCrop(sz, sz),\n            #A.Transpose(p=0.5),\n            A.HorizontalFlip(p=0.5),\n            #A.VerticalFlip(p=0.5),\n            #A.ShiftScaleRotate(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, p=1.0),\n            ToTensorV2(p=1.0),\n        ], p=1.)\n\n\ndef valid_transforms():\n    return Compose([\n            A.Resize(sz, sz),\n            A.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 (modified to support SnapMix)"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaNet(nn.Module):\n    def __init__(self):\n        super().__init__()\n        backbone = timm.create_model(TIMM_MODEL, pretrained=True)\n        n_features = backbone.fc.in_features\n        self.backbone = nn.Sequential(*backbone.children())[:-2]\n        self.classifier = nn.Linear(n_features, 5)\n        self.pool = nn.AdaptiveAvgPool2d((1, 1))\n\n    def forward_features(self, x):\n        x = self.backbone(x)\n        return x\n\n    def forward(self, x):\n        feats = self.forward_features(x)\n        x = self.pool(feats).view(x.size(0), -1)\n        x = self.classifier(x)\n        return x, feats","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Metrics"},{"metadata":{"trusted":true},"cell_type":"code","source":"def accuracy_metric(input, targs):\n    return accuracy_score(targs.cpu(), input.cpu())\n\ndef print_scores(scores):\n    kaggle_metric = np.average(scores)\n    print(\"Kaggle Metric: %f\" % (kaggle_metric))\n    \n    return kaggle_metric","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Checkpoint method"},{"metadata":{"trusted":true},"cell_type":"code","source":"def checkpoint(model, optimizer, epoch, current_metric, best_metric, fold):\n    print(\"Metric improved from %f to %f , Saving Model at Epoch #%d\" % (best_metric, current_metric, epoch))\n    ckpt = {\n        'model': CassavaNet(),\n        'state_dict': model.state_dict(),\n        #'optimizer' : optimizer.state_dict(),  # Commenting this out to cheap out on space\n        'metric': current_metric\n    }\n    torch.save(ckpt, 'ckpt_%s-%d-%d.pth' % (TIMM_MODEL, sz, fold))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Create folds"},{"metadata":{"trusted":true},"cell_type":"code","source":"folds = StratifiedKFold(n_splits=NUM_FOLDS, shuffle=True, \n                        random_state=SEED).split(np.arange(train_df.shape[0]), \n                                                 train_df.label.values)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# SnapMix Augmentation"},{"metadata":{"trusted":true},"cell_type":"code","source":"def rand_bbox(size, lam):\n    W = size[2]\n    H = size[3]\n    cut_rat = np.sqrt(1. - lam)\n    cut_w = np.int(W * cut_rat)\n    cut_h = np.int(H * cut_rat)\n\n    # uniform\n    cx = np.random.randint(W)\n    cy = np.random.randint(H)\n\n    bbx1 = np.clip(cx - cut_w // 2, 0, W)\n    bby1 = np.clip(cy - cut_h // 2, 0, H)\n    bbx2 = np.clip(cx + cut_w // 2, 0, W)\n    bby2 = np.clip(cy + cut_h // 2, 0, H)\n\n    return bbx1, bby1, bbx2, bby2\n\ndef get_spm(input,target,model):\n    imgsize = (sz, sz)\n    bs = input.size(0)\n    with torch.no_grad():\n        output,fms = model(input)\n        clsw = model.classifier\n        weight = clsw.weight.data\n        bias = clsw.bias.data\n        weight = weight.view(weight.size(0),weight.size(1),1,1)\n        fms = F.relu(fms)\n        poolfea = F.adaptive_avg_pool2d(fms,(1,1)).squeeze()\n        clslogit = F.softmax(clsw.forward(poolfea))\n        logitlist = []\n        for i in range(bs):\n            logitlist.append(clslogit[i,target[i]])\n        clslogit = torch.stack(logitlist)\n\n        out = F.conv2d(fms, weight, bias=bias)\n\n        outmaps = []\n        for i in range(bs):\n            evimap = out[i,target[i]]\n            outmaps.append(evimap)\n\n        outmaps = torch.stack(outmaps)\n        if imgsize is not None:\n            outmaps = outmaps.view(outmaps.size(0),1,outmaps.size(1),outmaps.size(2))\n            outmaps = F.interpolate(outmaps,imgsize,mode='bilinear',align_corners=False)\n\n        outmaps = outmaps.squeeze()\n\n        for i in range(bs):\n            outmaps[i] -= outmaps[i].min()\n            outmaps[i] /= outmaps[i].sum()\n\n\n    return outmaps,clslogit\n\n\ndef snapmix(input, target, alpha, model=None):\n\n    r = np.random.rand(1)\n    lam_a = torch.ones(input.size(0))\n    lam_b = 1 - lam_a\n    target_b = target.clone()\n\n    if True:\n        wfmaps,_ = get_spm(input, target, model)\n        bs = input.size(0)\n        lam = np.random.beta(alpha, alpha)\n        lam1 = np.random.beta(alpha, alpha)\n        rand_index = torch.randperm(bs).cuda()\n        wfmaps_b = wfmaps[rand_index,:,:]\n        target_b = target[rand_index]\n\n        same_label = target == target_b\n        bbx1, bby1, bbx2, bby2 = rand_bbox(input.size(), lam)\n        bbx1_1, bby1_1, bbx2_1, bby2_1 = rand_bbox(input.size(), lam1)\n\n        area = (bby2-bby1)*(bbx2-bbx1)\n        area1 = (bby2_1-bby1_1)*(bbx2_1-bbx1_1)\n\n        if  area1 > 0 and  area>0:\n            ncont = input[rand_index, :, bbx1_1:bbx2_1, bby1_1:bby2_1].clone()\n            ncont = F.interpolate(ncont, size=(bbx2-bbx1,bby2-bby1), mode='bilinear', align_corners=True)\n            input[:, :, bbx1:bbx2, bby1:bby2] = ncont\n            lam_a = 1 - wfmaps[:,bbx1:bbx2,bby1:bby2].sum(2).sum(1)/(wfmaps.sum(2).sum(1)+1e-8)\n            lam_b = wfmaps_b[:,bbx1_1:bbx2_1,bby1_1:bby2_1].sum(2).sum(1)/(wfmaps_b.sum(2).sum(1)+1e-8)\n            tmp = lam_a.clone()\n            lam_a[same_label] += lam_b[same_label]\n            lam_b[same_label] += tmp[same_label]\n            lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (input.size()[-1] * input.size()[-2]))\n            lam_a[torch.isnan(lam_a)] = lam\n            lam_b[torch.isnan(lam_b)] = 1-lam\n\n    return input,target,target_b,lam_a.cuda(),lam_b.cuda()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# SnapMix Criterion (Loss)"},{"metadata":{"trusted":true},"cell_type":"code","source":"class SnapMixLoss(nn.Module):\n    def __init__(self):\n        super().__init__()\n        \n    def forward(self, criterion, outputs, ya, yb, lam_a, lam_b):\n        loss_a = criterion(outputs, ya)\n        loss_b = criterion(outputs, yb)\n        loss = torch.mean(loss_a * lam_a + loss_b * lam_b)\n        return loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train & Validate"},{"metadata":{"trusted":true},"cell_type":"code","source":"for fold_num, (train_split, valid_split) in enumerate(folds):\n    train_set = train_df.iloc[train_split].reset_index(drop=True)\n    valid_set = train_df.iloc[valid_split].reset_index(drop=True)\n    \n    train_ds = CassavaDataset(dataframe=train_set,\n                          root_dir=DATA_PATH + 'train_images',\n                          transforms=train_transforms())\n    \n    valid_ds = CassavaDataset(dataframe=valid_set,\n                          root_dir=DATA_PATH + 'train_images',\n                          transforms=valid_transforms())\n    \n    train_dl = torch.utils.data.DataLoader(train_ds, batch_size=bs, \n                                           shuffle=True, num_workers=8, drop_last=True,\n                                           pin_memory=True)\n    valid_dl = torch.utils.data.DataLoader(valid_ds, batch_size=bs, \n                                           shuffle=False, num_workers=8, \n                                           pin_memory=True)\n    \n    losses = []\n    batches = len(train_dl)\n    val_batches = len(valid_dl)\n    best_metric = 0\n    \n    model = CassavaNet().to(device)\n    criterion = nn.CrossEntropyLoss(reduction='none').to(device)\n    val_criterion = nn.CrossEntropyLoss().to(device)\n    snapmix_criterion = SnapMixLoss().to(device)\n    param_groups = [\n       {'params': model.backbone.parameters(), 'lr': 1e-2},\n       {'params': model.classifier.parameters()},\n    ]\n    optimizer = torch.optim.SGD(param_groups, lr=1e-1, momentum=0.9,\n                                weight_decay=1e-4, nesterov=True)\n    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[1,20,40], \n                                                     gamma=0.1, last_epoch=-1, verbose=True)\n    scaler = GradScaler()\n    \n    for epoch in range(EPOCHS):\n        # ----------------- TRAINING  ----------------- \n        train_loss = 0\n        progress = tqdm(enumerate(train_dl), desc=\"Loss: \", total=batches)\n\n        model.train()\n        for i, data in progress:\n            image, label = data.values()\n            X, y = image.to(device).float(), label.to(device).long()\n            \n            with autocast():\n                \n                rand = np.random.rand()\n                if rand > (1.0-SNAPMIX_PCT):\n                    X, ya, yb, lam_a, lam_b = snapmix(X, y, SNAPMIX_ALPHA, model)\n                    outputs, _ = model(X)\n                    loss = snapmix_criterion(criterion, outputs, ya, yb, lam_a, lam_b)\n                else:\n                    outputs, _ = model(X)\n                    loss = torch.mean(criterion(outputs, y))\n                \n            scaler.scale(loss).backward()\n            # Accumulate gradients\n            if ((i + 1) % GRAD_ACCUM_STEPS == 0) or ((i + 1) == len(train_dl)):\n                scaler.step(optimizer)\n                scaler.update()\n                optimizer.zero_grad()\n\n            train_loss += loss.item()\n            cur_step = i+1\n            trn_epoch_result = dict()\n            trn_epoch_result['Epoch'] = epoch + 1\n            trn_epoch_result['train_loss'] = round(train_loss/cur_step, 4)\n\n            progress.set_description(str(trn_epoch_result))\n\n        scheduler.step()\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n\n        # ----------------- VALIDATION  ----------------- \n        val_loss = 0\n        scores = []\n\n        model.eval()\n        with torch.no_grad():\n            for i, data in enumerate(valid_dl):\n                image, label = data.values()\n                X, y = image.to(device), label.to(device)\n                outputs, _ = model(X)\n                l = val_criterion(outputs, y)\n                val_loss += l.item()\n\n                preds = F.softmax(outputs).argmax(axis=1)\n                scores.append(accuracy_metric(preds, y))\n\n        epoch_result = dict()\n        epoch_result['Epoch'] = epoch + 1\n        epoch_result['train_loss'] = round(train_loss/batches, 4)\n        epoch_result['val_loss'] = round(val_loss/val_batches, 4)\n\n        print(epoch_result)\n\n        # Check if we need to save\n        current_metric = print_scores(scores)\n        if current_metric > best_metric:\n            checkpoint(model, optimizer, epoch+1, current_metric, best_metric, fold_num)\n            best_metric = current_metric\n            \n    del model, optimizer, train_dl, valid_dl, scaler, scheduler\n    torch.cuda.empty_cache()\n    \n    # Train only a single fold\n    break","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"*Please Upvote if you liked the kernel ! Cheers.*"}],"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}