{"cells":[{"metadata":{},"cell_type":"markdown","source":"## About this kernel  \n#### Used SEresnext50 with SnapMix augmentation\n## Source kernel  \n#### This notebook was written by refering these great kernels below, so please don't forget to check and upvote them.  \n* https://www.kaggle.com/sachinprabhu/pytorch-resnet50-snapmix-train-pipeline  \n* https://www.kaggle.com/shaolihuang/training-with-snapmix  "},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"!pip install timm\n!pip install torchsummary","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import cv2\nimport torch\nfrom torch import nn\nimport random\nimport os\nimport time\nfrom tqdm import tqdm_notebook as tqdm\nimport torchvision\nfrom torchvision import transforms\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda.amp import autocast, GradScaler\nimport torch.nn.functional as F\nimport pickle\nimport timm\n\nfrom PIL import Image\nfrom albumentations import *\nfrom albumentations.pytorch import ToTensorV2","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Config"},{"metadata":{"trusted":true},"cell_type":"code","source":"CFG = {\n    'scheduler' : 'MultiStepLR',\n    'optimizer' : 'SGD',\n    'snapmix_alpha' : 5.0,\n    'snapmix_ptc' : 1.0, \n    'model_name' : 'seresnext50_32x4d',\n    'num_classes' : 5, \n    'accum_iter' : 1,\n    'current_fold': 0,\n    'fold_num' : 5, \n    'img_size': 448,\n    'resize' : 512,\n    'epochs': 10,\n    'milestones': [1, 10, 20], # MultiStepLR\n    'train_bs': 32,\n    'valid_bs': 64,\n    'weight_decay_Adam' : 1e-6,\n    'weight_decay_SGD' : 1e-2,\n    'num_workers': 8,\n    'device': 'cuda:0',\n    'seed' : 3,\n    'verbose_step': 1,\n    'T_0': 10, # CosineAnnealingWarmRestarts\n    'lr': 1e-3,\n    'min_lr': 1e-5 # CosineAnnealingWarmRestarts\n}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Load csv file"},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv('../input/colabcassava/train.csv')\ntrain.label.value_counts()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed = 3):\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    \ndef get_img(path):\n    im_bgr = cv2.imread(path)\n    im_rgb = im_bgr[:, :, ::-1]\n    #print(im_rgb)\n    return im_rgb","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 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        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":{},"cell_type":"markdown","source":"### Rand augmentation"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transforms():\n    return Compose([\n            Resize(CFG['resize'], CFG['resize'], p = 1.),\n            RandomCrop(CFG['img_size'], CFG['img_size'], p = 1.),\n            HorizontalFlip(p=0.5),\n            #VerticalFlip(p=0.5),\n            #Transpose(p = 0.5),\n            #Rotate(limit = 10, 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            ToTensorV2(p=1.0),\n        ], p=1.)\n    \ndef get_valid_transforms():\n    return Compose([\n            Resize(CFG['resize'], CFG['resize'], p = 1.),\n            CenterCrop(CFG['img_size'], CFG['img_size'], p=1.),\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":{"trusted":true},"cell_type":"code","source":"print(\"Available SEresnext Models: \")\ntimm.list_models(\"seresnext*\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaNet(nn.Module):\n    def __init__ (self):\n        super().__init__()\n        backbone = timm.create_model(CFG['model_name'], pretrained=True)\n        n_features = backbone.fc.in_features\n        self.backbone = nn.Sequential(*backbone.children())[:-2]\n        self.classifier = nn.Linear(n_features, CFG['num_classes'])\n        self.pool = nn.AdaptiveAvgPool2d((1, 1))\n        #self.model.fc = nn.Linear(n_features, CFG['num_classes'])\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":"### summary model"},{"metadata":{"trusted":true},"cell_type":"code","source":"from torchsummary import summary\ndevice = torch.device(CFG['device'])\nmodel = CassavaNet().to(device)\nprint(summary(model, (3, CFG['img_size'], CFG['img_size'])))\ndel model","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 = (CFG['img_size'], CFG['img_size'])\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 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":{"trusted":true},"cell_type":"code","source":"def prepare_dataloader(df, trn_idx, val_idx, \n                       data_root=None):\n    \n    train_ = df.loc[trn_idx,:].reset_index(drop=True)\n    valid_ = df.loc[val_idx,:].reset_index(drop=True)\n        \n    train_ds = CassavaDataset(train_, data_root, transforms=get_train_transforms())\n    valid_ds = CassavaDataset(valid_, data_root, transforms=get_valid_transforms())\n    \n    train_loader = torch.utils.data.DataLoader(\n        train_ds,\n        batch_size=CFG['train_bs'],\n        pin_memory=False,\n        drop_last=False,\n        shuffle=True,        \n        num_workers=CFG['num_workers'],\n    )\n    val_loader = torch.utils.data.DataLoader(\n        valid_ds, \n        batch_size=CFG['valid_bs'],\n        num_workers=CFG['num_workers'],\n        shuffle=False,\n        pin_memory=False,\n    )\n    return train_loader, val_loader","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_one_epoch(epoch, model, loss_fn, optimizer, train_loader, device):\n    model.train()\n\n    t = time.time()\n    running_loss = None\n    sample = 0\n\n    pbar = tqdm(enumerate(train_loader), total=len(train_loader))\n    for step, data in pbar:\n        (imgs, image_labels) = data.values()\n        imgs = imgs.to(device).float()\n        image_labels = image_labels.to(device).long()\n\n        with autocast(): \n\n            outputs = model(imgs)\n            loss = loss_fn(outputs, image_labels)\n            \n            scaler.scale(loss).backward()\n\n            if running_loss is None:\n                running_loss = loss.item()*image_labels.shape[0]\n            else:\n                running_loss += loss.item()*image_labels.shape[0]\n\n            sample += image_labels.shape[0]\n\n            if ((step + 1) %  CFG['accum_iter'] == 0) or ((step + 1) == len(train_loader)):\n                # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)\n\n                scaler.step(optimizer)\n                scaler.update()\n                optimizer.zero_grad() \n\n            if ((step + 1) % CFG['verbose_step'] == 0) or ((step + 1) == len(train_loader)):\n                description = f'epoch {epoch} loss: {running_loss/sample:.4f}'\n                \n                pbar.set_description(description)\n\ndef train_one_epoch_snapmix(epoch, model, loss_fn, optimizer, train_loader, device):\n    model.train()\n    snapmix_criterion = SnapMixLoss().to(device)\n    t = time.time()\n    running_loss = None\n    sample = 0\n    count_snapmix = 0\n    pbar = tqdm(enumerate(train_loader), total=len(train_loader))\n    for step, data in pbar:\n        (imgs, image_labels) = data.values()\n        imgs = imgs.to(device).float()\n        image_labels = image_labels.to(device).long()\n\n        with autocast(): \n            \n            rand = np.random.rand()\n            loss = None\n            if rand > (1.0 - CFG['snapmix_ptc']):\n                count_snapmix += 1\n                imgs, ya, yb, lam_a, lam_b = snapmix(imgs, image_labels, CFG['snapmix_alpha'], model)\n                outputs, _ = model(imgs)\n                loss = snapmix_criterion(loss_fn, outputs, ya, yb, lam_a, lam_b)\n            else:\n                outputs, _ = model(imgs)\n                loss = torch.mean(loss_fn(outputs, image_labels))\n\n            #outputs = model(imgs)\n            #loss = loss_fn(outputs, image_labels)\n            \n            scaler.scale(loss).backward()\n\n            if running_loss is None:\n                running_loss = loss.item()*image_labels.shape[0]\n            else:\n                running_loss += loss.item()*image_labels.shape[0]\n\n            sample += image_labels.shape[0]\n\n            if ((step + 1) %  CFG['accum_iter'] == 0) or ((step + 1) == len(train_loader)):\n                # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)\n\n                scaler.step(optimizer)\n                scaler.update()\n                optimizer.zero_grad() \n\n            if ((step + 1) % CFG['verbose_step'] == 0) or ((step + 1) == len(train_loader)):\n                description = f'epoch {epoch} loss: {running_loss/sample:.4f}'\n                \n                pbar.set_description(description)\n    print('Number of snapmix iterations: {}/{}'.format(count_snapmix, len(train_loader)))\n\ndef valid_one_epoch(epoch, model, loss_fn, val_loader, device):\n    model.eval()\n\n    t = time.time()\n    loss_sum = 0\n    sample_num = 0\n    image_preds_all = []\n    image_targets_all = []\n    \n    pbar = tqdm(enumerate(val_loader), total=len(val_loader))\n    for step, data in pbar:\n        (imgs, image_labels) = data.values()\n        imgs = imgs.to(device).float()\n        image_labels = image_labels.to(device).long()\n        \n        image_preds, _ = model(imgs)\n        image_preds_all += [torch.argmax(image_preds, 1).detach().cpu().numpy()]\n        image_targets_all += [image_labels.detach().cpu().numpy()]\n        \n        loss = loss_fn(image_preds, image_labels)\n        \n        loss_sum += loss.item()*image_labels.shape[0]\n        sample_num += image_labels.shape[0]  \n\n        if ((step + 1) % CFG['verbose_step'] == 0) or ((step + 1) == len(val_loader)):\n            description = f'epoch {epoch} loss: {loss_sum/sample_num:.4f}'\n            pbar.set_description(description)\n    \n    image_preds_all = np.concatenate(image_preds_all)\n    image_targets_all = np.concatenate(image_targets_all)\n    accuracy = (image_preds_all==image_targets_all).mean()\n    print('validation multi-class accuracy = {:.4f}'.format(accuracy))\n    \n    return accuracy","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Train"},{"metadata":{"trusted":true},"cell_type":"code","source":"if __name__ == '__main__':\n     # for training only, need nightly build pytorch\n\n    seed_everything(CFG['seed'])\n    \n    for fold in range(CFG['fold_num']):\n\n        if fold != CFG['current_fold']:\n            continue\n            \n        pickle_in = open(\"../input/5folds/train_fold \" + str(fold) + \".pickle\",\"rb\")\n        train_folds = pickle.load(pickle_in)\n        pickle_in.close()\n    \n        pickle_in = open(\"../input/5folds/valid_fold \" + str(fold) + \".pickle\",\"rb\")\n        valid_folds = pickle.load(pickle_in)\n        pickle_in.close()\n        \n        trn_idx = list(train_folds)\n        val_idx = list(valid_folds)\n        \n        print('Training with {} started'.format(fold))\n        print(len(trn_idx), len(val_idx))\n\n        train_loader, val_loader = prepare_dataloader(train, trn_idx, val_idx, \n                                    data_root='../input/colabcassava/handle_data')\n\n        device = torch.device(CFG['device'])\n        \n        model = CassavaNet().to(device)\n\n        scaler = GradScaler()   \n        param_groups = [\n            {'params': model.backbone.parameters(), 'lr': 1e-2},\n            {'params': model.classifier.parameters()},\n        ]\n\n        optimizer = None\n        if CFG['optimizer'] == 'SGD':\n            optimizer = torch.optim.SGD(param_groups, lr=1e-1, momentum=0.9,\n                                    weight_decay=CFG['weight_decay_SGD'], nesterov=True)\n        elif optimizer == 'Adam':\n            optimizer = torch.optim.Adam(param_groups, lr=CFG['lr'], \n                                         weight_decay=CFG['weight_decay_Adam'])\n\n        scheduler = None\n        if CFG['scheduler'] == 'CosineAnnealingWarmRestarts':\n            scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=CFG['T_0'], \n                                    T_mult=1, eta_min=CFG['min_lr'], last_epoch=-1, verbose = 1)\n        elif CFG['scheduler'] == 'ReduceLROnPlateau':\n            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.3, \n                                      patience=2, verbose=True, eps=1e-6)\n        elif CFG['scheduler'] == 'MultiStepLR':\n            scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=CFG['milestones'], \n                                                     gamma=0.1, last_epoch=-1, verbose=True)\n        \n        loss_tr = nn.CrossEntropyLoss(reduction='none').to(device)\n        loss_fn = nn.CrossEntropyLoss().to(device)\n        \n        best_accuracy = 0\n        #model.load_state_dict( \n        #           torch.load('/content/gdrive/MyDrive/Data/resnext50/resnext50_32x4d_fold_0.pth'))\n        for epoch in range(CFG['epochs']):\n            train_one_epoch_snapmix(epoch, model, loss_tr, optimizer, train_loader, device)\n            \n            valid_acc = 0\n            with torch.no_grad():\n                valid_acc = valid_one_epoch(epoch, model, loss_fn, val_loader, device)\n            if CFG['scheduler'] == 'ReduceLROnPlateau':\n                scheduler.step(valid_acc)\n            else:\n                scheduler.step()\n            if valid_acc > best_accuracy:\n                best_accuracy = valid_acc\n                print('Save best model at epoch {} : {}'.format(epoch, valid_acc))\n                torch.save(model.state_dict(), '{}_fold_{}.pth'.format(CFG['model_name'], fold))\n            torch.save(model.state_dict(), '{}_cur_fold_{}.pth'.format(CFG['model_name'], fold))\n        del model, optimizer, train_loader, val_loader, scaler, scheduler\n        torch.cuda.empty_cache()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"*If you see this kernel helpfully, please Upvote it! Have fun*"}],"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}