{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!git clone https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer.git\n!pip install adamp","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"import sys\nsys.path.append('../input/pytorch-image-models/pytorch-image-models-master')\nsys.path.append('./Ranger-Deep-Learning-Optimizer/ranger')\nsys.path.append('../input/image-fmix/FMix-master')\n# sys.setrecursionlimit(10**6)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom datetime import datetime\nfrom tqdm.notebook import tqdm\nfrom pprint import pprint\nimport cv2, glob, time, random, os\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\nimport timm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.transforms as T\nfrom torch.utils.data import Dataset, DataLoader\n# https://nvlabs.github.io/iccv2019-mixed-precision-tutorial/files/dusan_stosic_intro_to_mixed_precision_training.pdf\n# https://analyticsindiamag.com/pytorch-mixed-precision-training/\n# https://pytorch.org/docs/stable/notes/amp_examples.html\nfrom torch.cuda.amp import autocast, GradScaler\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau, CosineAnnealingWarmRestarts\nfrom torch.optim import Adam, AdamW, SGD\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom fmix import sample_mask\n\nfrom adamp import AdamP\n# https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer\nfrom ranger import Ranger  # this is from ranger.py\nfrom ranger913A import RangerVA  # this is from ranger913A.py\nfrom rangerqh import RangerQH  # this is from rangerqh.py\n\nfrom sklearn import preprocessing\nfrom sklearn.metrics import accuracy_score\nfrom sklearn.model_selection import StratifiedKFold","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# get the list of pretrained models\nmodel_names = timm.list_models()\npprint(model_names)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<a id = \"cont\"></a>\n## CFG"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"BATCH_SIZE = 16 # 8 for bigger architectures\nVAL_BATCH_SIZE = 8\nEPOCHS = 2 # train upto 10 epochs\nIMG_SIZE = 512 # 384 for bigger architectures\nITER_FREQ = 500\nNUM_WORKERS = 8\nMEAN = [0.485, 0.456, 0.406]\nSTD = [0.229, 0.224, 0.225]\nSEED = 1111\nN_FOLDS = 5\nTR_FOLDS = [0,1,2,3,4]\nSTART_FOLD = 0\n\nMODEL_PATH = None\n\nLR = 1e-4\nMIN_LR = 5e-5 # SAM, CosineAnnealingWarmRestarts\nWEIGHT_DECAY = 1e-6\nMOMENTUM = 0.9\nT_0 = EPOCHS # SAM, CosineAnnealingWarmRestarts\nMAX_NORM = 1000\n\nSNAPMIX = False\nSNAPMIX_ALPHA = 5.0\nSNAPMIX_PCT = 0.5\n\nMODEL_ARCH = 'tf_efficientnet_b3_ns' # tf_efficientnet_b4_ns, tf_efficientnet_b5_ns, resnext50_32x4d\nITERS_TO_ACCUMULATE = 1\n\nCUTMIX = False\nCM_START = 0\nCM_ALPHA = 1\nDECAY_POWER = 5\n\nBASE_OPTIMIZER = SGD #for SAM, Ranger\nOPTIMIZER = 'AdamW' # Ranger, Adam, AdamP, SGD, SAM\n\nSCHEDULER = 'CosineAnnealingWarmRestarts' # ReduceLROnPlateau, CosineAnnealingLR, CosineAnnealingWarmRestarts, OneCycleLR\nSCHEDULER_UPDATE = 'epoch' # batch\n\nCRITERION = 'TaylorCrossEntropyLoss' # CrossEntropyLoss, TaylorSmoothedLoss, LabelSmoothedLoss\nLABEL_SMOOTH = 0.5\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"class AverageMeter(object):\n    \n    # Keeps track of most recent, average, sum, and count of a metric.\n    \n    def __init__(self):\n        self.reset()\n        \n    def reset(self):\n        self.val = 0\n        self.sum = 0\n        self.avg = 0\n        self.count = 0\n        \n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val*n\n        self.count += n\n        self.avg = self.sum / self.count\n\ndef seed_torch(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    \nseed_torch(SEED)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"TRAIN_DIR = '../input/cassava-leaf-disease-merged/train/'\ndf = pd.read_csv('../input/cassava-leaf-disease-merged/merged.csv')\n# df","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"folds = df.copy()\nFold = StratifiedKFold(n_splits=N_FOLDS, shuffle=True, random_state=SEED)\nfor n, (train_index, val_index) in enumerate(Fold.split(folds, folds['label'])):\n    folds.loc[val_index, 'fold'] = int(n)\nfolds['fold'] = folds['fold'].astype(int)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"class CLD_train_DS(Dataset):\n    \n    def __init__(self, df, transform=None):\n        self.df = df\n        self.image_id = df['image_id'].values\n        self.label = df['label'].values\n        self.transform = transform\n        \n    def __len__(self) -> int:\n        return len(self.df)\n        \n    def __getitem__(self, idx):\n        \n        image_id = self.image_id[idx]\n        img = cv2.imread(TRAIN_DIR + image_id)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32)\n        img = img / 255.0\n        \n        if self.transform:\n            trans_img = self.transform(image=img)\n            img = trans_img['image']\n            \n        label = torch.tensor(self.label[idx]).long()\n        \n        return img, label","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"def get_transform(*, train=True):\n    \n    if train:\n        return A.Compose([\n            A.RandomResizedCrop(IMG_SIZE, IMG_SIZE),\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=MEAN, std=STD),\n            ToTensorV2(),\n        ])\n    else:\n        return A.Compose([\n#             A.CenterCrop(IMG_SIZE, IMG_SIZE),\n            A.Resize(IMG_SIZE, IMG_SIZE),\n            A.Normalize(mean=MEAN, std=STD, max_pixel_value=255.0, p=1.0),\n            ToTensorV2(),\n        ])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"[SAM optimizer](https://github.com/davda54/sam)"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"class SAM(torch.optim.Optimizer):\n    def __init__(self, params, base_optimizer, rho=0.05, **kwargs):\n        assert rho >= 0.0, f\"Invalid rho, should be non-negative: {rho}\"\n\n        defaults = dict(rho=rho, **kwargs)\n        super(SAM, self).__init__(params, defaults)\n\n        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)\n        self.param_groups = self.base_optimizer.param_groups\n\n    @torch.no_grad()\n    def first_step(self, zero_grad=False):\n        grad_norm = self._grad_norm()\n        for group in self.param_groups:\n            scale = group[\"rho\"] / (grad_norm + 1e-12)\n\n            for p in group[\"params\"]:\n                if p.grad is None: continue\n                e_w = p.grad * scale.to(p)\n                p.add_(e_w)  # climb to the local maximum \"w + e(w)\"\n                self.state[p][\"e_w\"] = e_w\n\n        if zero_grad: self.zero_grad()\n\n    @torch.no_grad()\n    def second_step(self, zero_grad=False):\n        for group in self.param_groups:\n            for p in group[\"params\"]:\n                if p.grad is None: continue\n                p.sub_(self.state[p][\"e_w\"])  # get back to \"w\" from \"w + e(w)\"\n\n        self.base_optimizer.step()  # do the actual \"sharpness-aware\" update\n\n        if zero_grad: self.zero_grad()\n\n    @torch.no_grad()\n    def step(self, closure=None):\n        assert closure is not None, \"Sharpness Aware Minimization requires closure, but it was not provided\"\n        closure = torch.enable_grad()(closure)  # the closure should do a full forward-backward pass\n\n        self.first_step(zero_grad=True)\n        closure()\n        self.second_step()\n\n    def _grad_norm(self):\n        shared_device = self.param_groups[0][\"params\"][0].device  # put everything on the same device, in case of model parallelism\n        norm = torch.norm(\n                    torch.stack([\n                        p.grad.norm(p=2).to(shared_device)\n                        for group in self.param_groups for p in group[\"params\"]\n                        if p.grad is not None\n                    ]),\n                    p=2\n               )\n        return norm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"[Back to CFG(Click here)](#cont)"},{"metadata":{},"cell_type":"markdown","source":"## Loss Functions"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"class LabelSmoothingLoss(nn.Module): \n    def __init__(self, classes=5, smoothing=0.0, dim=-1): \n        super(LabelSmoothingLoss, self).__init__() \n        self.confidence = 1.0 - smoothing \n        self.smoothing = smoothing \n        self.cls = classes \n        self.dim = dim \n    def forward(self, pred, target): \n        if CRITERION == 'LabelSmoothingLoss':\n            pred = pred.log_softmax(dim=self.dim) \n        with torch.no_grad():\n            true_dist = torch.zeros_like(pred) \n            true_dist.fill_(self.smoothing / (self.cls - 1)) \n            true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence) \n        return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))\n    \nclass TaylorSoftmax(nn.Module):\n    '''\n    This is the autograd version\n    '''\n    def __init__(self, dim=1, n=2):\n        super(TaylorSoftmax, self).__init__()\n        assert n % 2 == 0\n        self.dim = dim\n        self.n = n\n\n    def forward(self, x):\n        '''\n        usage similar to nn.Softmax:\n            >>> mod = TaylorSoftmax(dim=1, n=4)\n            >>> inten = torch.randn(1, 32, 64, 64)\n            >>> out = mod(inten)\n        '''\n        fn = torch.ones_like(x)\n        denor = 1.\n        for i in range(1, self.n+1):\n            denor *= i\n            fn = fn + x.pow(i) / denor\n        out = fn / fn.sum(dim=self.dim, keepdims=True)\n        return out\n    \nclass TaylorCrossEntropyLoss(nn.Module):\n    '''\n    This is the autograd version\n    '''\n    def __init__(self, n=2, ignore_index=-1, reduction='mean'):\n        super(TaylorCrossEntropyLoss, self).__init__()\n        assert n % 2 == 0\n        self.taylor_softmax = TaylorSoftmax(dim=1, n=n)\n        self.reduction = reduction\n        self.ignore_index = ignore_index\n\n    def forward(self, logits, labels):\n        '''\n        usage similar to nn.CrossEntropyLoss:\n            >>> crit = TaylorCrossEntropyLoss(n=4)\n            >>> inten = torch.randn(1, 10, 64, 64)\n            >>> label = torch.randint(0, 10, (1, 64, 64))\n            >>> out = crit(inten, label)\n        '''\n        log_probs = self.taylor_softmax(logits).log()\n        loss = F.nll_loss(log_probs, labels, reduction=self.reduction,\n                ignore_index=self.ignore_index)\n        return loss\n    \nclass TaylorSmoothedLoss(nn.Module):\n\n    def __init__(self, n=2, ignore_index=-1, reduction='mean', smoothing=0.2):\n        super(TaylorSmoothedLoss, self).__init__()\n        assert n % 2 == 0\n        self.taylor_softmax = TaylorSoftmax(dim=1, n=n)\n        self.reduction = reduction\n        self.ignore_index = ignore_index\n        self.lab_smooth = LabelSmoothingLoss(5, smoothing=LABEL_SMOOTH)\n\n    def forward(self, logits, labels):\n\n        log_probs = self.taylor_softmax(logits).log()\n        #loss = F.nll_loss(log_probs, labels, reduction=self.reduction,\n        #        ignore_index=self.ignore_index)\n        loss = self.lab_smooth(log_probs, labels)\n        return loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Augmentations"},{"metadata":{},"cell_type":"markdown","source":"[SnapMix](https://github.com/Shaoli-Huang/SnapMix)"},{"metadata":{"trusted":true,"scrolled":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 = (IMG_SIZE, 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()\n\nclass 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":"CutMix"},{"metadata":{"trusted":true},"cell_type":"code","source":"def cutmix(batch, alpha):\n    data, targets = batch\n\n    indices = torch.randperm(data.size(0))\n    shuffled_data = data[indices]\n    shuffled_targets = targets[indices]\n\n    lam = np.random.beta(alpha, alpha)\n\n    image_h, image_w = data.shape[2:]\n    cx = np.random.uniform(0, image_w)\n    cy = np.random.uniform(0, image_h)\n    w = image_w * np.sqrt(1 - lam)\n    h = image_h * np.sqrt(1 - lam)\n    x0 = int(np.round(max(cx - w / 2, 0)))\n    x1 = int(np.round(min(cx + w / 2, image_w)))\n    y0 = int(np.round(max(cy - h / 2, 0)))\n    y1 = int(np.round(min(cy + h / 2, image_h)))\n\n    data[:, :, y0:y1, x0:x1] = shuffled_data[:, :, y0:y1, x0:x1]\n    return_targets = torch.zeros((len(targets),3),dtype=torch.int64)\n    return_targets[:,0] = targets\n    return_targets[:,1] = shuffled_targets\n    return_targets[0,2] = lam\n\n    return data, return_targets\n        \nclass CutMixCollator:\n    def __init__(self, alpha):\n        self.alpha = alpha\n\n    def __call__(self, batch):\n        batch = torch.utils.data.dataloader.default_collate(batch)\n        batch = cutmix(batch, self.alpha)\n        return batch\n    \nclass CutMixCriterion(nn.Module):\n    def __init__(self, criterion):\n        super(CutMixCriterion, self).__init__()\n        self.criterion = criterion\n\n    def forward(self, preds, labels):\n        targets1 = labels[:,0] \n        targets2 = labels[:,1]\n        lam = labels[0,2]\n        loss = lam * self.criterion(preds, targets1) + (1 - lam) * self.criterion(preds, targets2)\n        return loss","execution_count":null,"outputs":[]},{"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    return bbx1, bby1, bbx2, bby2\n\ndef cutmix_(data, target, alpha):\n    indices = torch.randperm(data.size(0))\n    shuffled_data = data[indices]\n    shuffled_target = target[indices]\n\n    lam = np.clip(np.random.beta(alpha, alpha),0.3,0.4)\n    bbx1, bby1, bbx2, bby2 = rand_bbox(data.size(), lam)\n    new_data = data.clone()\n    new_data[:, :, bby1:bby2, bbx1:bbx2] = data[indices, :, bby1:bby2, bbx1:bbx2]\n    # adjust lambda to exactly match pixel ratio\n    lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (data.size()[-1] * data.size()[-2]))\n    targets = (target, shuffled_target, lam)\n\n    return new_data, targets\n\ndef fmix(data, targets, alpha, decay_power, shape, max_soft=0.0, reformulate=False):\n    lam, mask = sample_mask(alpha, decay_power, shape, max_soft, reformulate)\n    #mask =torch.tensor(mask, device=device).float()\n    indices = torch.randperm(data.size(0))\n    shuffled_data = data[indices]\n    shuffled_targets = targets[indices]\n    x1 = torch.from_numpy(mask).to(DEVICE)*data\n    x2 = torch.from_numpy(1-mask).to(DEVICE)*shuffled_data\n    targets=(targets, shuffled_targets, lam)\n    \n    return (x1+x2), targets","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"# Modified for SnapMix\nclass CassavaNet(nn.Module):\n    def __init__(self):\n        super().__init__()\n        backbone = timm.create_model(MODEL_ARCH, pretrained=True)\n        n_features = backbone.classifier.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\n\nclass CustomEffNet(nn.Module):\n    def __init__(self, model_arch, n_class, pretrained=True):\n        super().__init__()\n        self.model = timm.create_model(model_arch, pretrained, n_class)\n        n_features = self.model.classifier.in_features\n        self.model.classifier = nn.Linear(n_features, n_class)\n        \n    def forward(self, x):\n        x = self.model(x)\n        return x\n    \nclass CustomResNext(nn.Module):\n    def __init__(self, model_arch, n_class, pretrained=False):\n        super().__init__()\n        self.model = timm.create_model(model_arch, pretrained=pretrained)\n        n_features = self.model.fc.in_features\n        self.model.classifier = nn.Linear(n_features, n_class)\n        \n    def forward(self, x):\n        x = self.model(x)\n        return x\n\nclass CustomViT(nn.Module):\n    def __init__(self, model_arch, n_class, pretrained=False):\n        super().__init__()\n        self.model = timm.create_model(model_arch, pretrained=pretrained)\n        n_features = self.model.head.in_features\n        self.model.head = nn.Linear(n_features, n_class)\n\n    def forward(self, x):\n        x = self.model(x)\n        return x\n\nclass CustomDeiT(nn.Module):\n    def __init__(self, model_arch, n_class, pretrained=False):\n        super().__init__()\n        self.model = torch.hub.load('facebookresearch/deit:main', model_arch, pretrained=pretrained)\n        n_features = self.model.head.in_features\n        self.model.head = nn.Linear(n_features, n_class)\n\n    def forward(self, x):\n        x = self.model(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"[Back to CFG(Click here)](#cont)"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"def GetCriterion(criterion_name, criterion=None):\n#     if criterion_name == 'BiTemperedLoss':\n#         criterion = BiTemperedLogistic()\n#     elif criterion_name == 'SymmetricCrossEntropyLoss':\n#         criterion = SymmetricCrossEntropy()\n    if criterion_name == 'CrossEntropyLoss':\n        criterion = nn.CrossEntropyLoss()\n    elif criterion_name == 'LabelSmoothingLoss':\n        criterion = LabelSmoothingLoss()\n#     elif criterion_name == 'FocalLoss':\n#         criterion = FocalLoss()\n#     elif criterion_name == 'FocalCosineLoss':\n#         criterion = FocalCosineLoss()\n    elif criterion_name == 'TaylorCrossEntropyLoss':\n        criterion = TaylorCrossEntropyLoss()\n    elif criterion_name == 'TaylorSmoothedLoss':\n        criterion = TaylorSmoothedLoss()\n    elif criterion_name == 'CutMix':\n        criterion = CutMixCriterion(criterion)\n    elif criterion_name == 'SnapMix':\n        criterion = SnapMixLoss()\n    return criterion\n    \n    \ndef GetScheduler(scheduler_name, optimizer, batches=None):\n    #['ReduceLROnPlateau', 'CosineAnnealingLR', 'CosineAnnealingWarmRestarts', 'OneCycleLR', 'GradualWarmupSchedulerV2']\n    if scheduler_name == 'OneCycleLR':\n        return torch.optim.lr_scheduler.OneCycleLR(optimizer,max_lr = 1e-2,epochs = CFG.EPOCHS,\n                                                   steps_per_epoch = batches+1,pct_start = 0.1)\n    if scheduler_name == 'CosineAnnealingWarmRestarts':\n        return torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0 = T_0, T_mult=1,\n                                                                    eta_min=MIN_LR, last_epoch=-1)\n    elif scheduler_name == 'CosineAnnealingLR':\n        return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max = CFG.T_max, eta_min=0, last_epoch=-1)\n    elif scheduler_name == 'ReduceLROnPlateau':\n        return torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.1, patience=1, threshold=0.0001,\n                                                          cooldown=0, min_lr=MIN_LR)\n#     elif scheduler_name == 'GradualWarmupSchedulerV2':\n#         return GradualWarmupSchedulerV2(optimizer=optimizer)\n    \ndef GetOptimizer(optimizer_name,parameters):\n    #['Adam','Ranger']\n    if optimizer_name == 'Adam':\n#         if CFG.scheduler_name == 'GradualWarmupSchedulerV2':\n#             return torch.optim.Adam(parameters, lr=CFG.LR_START, weight_decay=CFG.weight_decay, amsgrad=False)\n#         else:\n        return torch.optim.Adam(parameters, lr=LR, weight_decay=WEIGHT_DECAY, amsgrad=False)\n    elif optimizer_name == 'AdamW':\n#         if CFG.scheduler_name == 'GradualWarmupSchedulerV2':\n#             return torch.optim.AdamW(parameters, lr=CFG.LR_START, weight_decay=CFG.weight_decay, amsgrad=False)\n#         else:\n        return torch.optim.Adam(parameters, lr=LR, weight_decay=WEIGHT_DECAY, amsgrad=False)\n    elif optimizer_name == 'AdamP':\n#         if CFG.scheduler_name == 'GradualWarmupSchedulerV2':\n#             return AdamP(parameters, lr=CFG.LR_START, weight_decay=CFG.weight_decay)\n#         else:\n        return AdamP(parameters, lr=LR, weight_decay=WEIGHT_DECAY)\n    elif optimizer_name == 'Ranger':\n        return Ranger(parameters, lr = LR, alpha = 0.5, k = 6, N_sma_threshhold = 5, \n                      betas = (0.95,0.999), weight_decay=WEIGHT_DECAY)\n    elif optimizer_name == 'SAM':\n        return SAM(parameters, BASE_OPTIMIZER, lr=0.1, momentum=0.9,weight_decay=0.0005)\n    \n    elif optimizer_name == 'AdamP':\n        return AdamP(parameters, lr=LR, weight_decay=WEIGHT_DECAY)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train and validation functions"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"def train_fn(model, dataloader, device, epoch, optimizer, criterion, scheduler):\n    \n    data_time = AverageMeter()\n    batch_time = AverageMeter()\n    losses = AverageMeter()\n    accuracies = AverageMeter()\n    \n    model.train()\n    # https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-accumulation\n    scaler = GradScaler()\n    start_time = time.time()\n    loader = tqdm(dataloader, total=len(dataloader))\n    for step, (images, labels) in enumerate(loader):\n        \n        images = images.to(device).float()\n        labels = labels.to(device)\n        data_time.update(time.time() - start_time)\n        \n#         mix_decision = np.random.rand()\n#         if mix_decision < 0.25:\n#             images, labels = cutmix_(images, labels, CM_ALPHA)\n#             images = images.float()\n#         elif mix_decision >=0.25 and mix_decision < 0.5:\n#             images, labels = fmix(images, labels, alpha=CM_ALPHA, decay_power=DECAY_POWER, shape=(IMG_SIZE,IMG_SIZE))\n#             images = images.float()\n\n        # Run forward pass with autocasting.\n        with autocast():\n            \n            if SNAPMIX:\n                rand = np.random.rand()\n                if rand > (1.0-SNAPMIX_PCT):\n                    X, ya, yb, lam_a, lam_b = snapmix(images, labels, SNAPMIX_ALPHA, model)\n                    output, _ = model(X)\n                    snapmix_criterion = GetCriterion('SnapMix')\n                    loss = snapmix_criterion(criterion, output, ya, yb, lam_a, lam_b)\n                else:\n                    output, _ = model(images)\n                    criterion = GetCriterion('CrossEntropyLoss').to(device)\n                    loss = criterion(output, labels)\n            else:\n                output = model(images)\n                loss = criterion(output, labels)\n#                 if mix_decision < 0.50:\n#                     loss = criterion(output, labels[0]) * labels[2] + criterion(output, labels[1]) * (1. - labels[2])\n#                 else:\n#                     loss = criterion(output, labels)\n                if CUTMIX:\n                    losses.update(loss, BATCH_SIZE)\n                else:\n                    losses.update(loss.item(), BATCH_SIZE)\n        \n        # Call .backward() to accumulate scaled gradients.\n        # Backward passes under autocast are not recommended.\n        # Backward ops run in the same dtype autocast chose for corresponding forward ops.\n        scaler.scale(loss).backward()\n        grad_norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm = MAX_NORM)\n\n        if not CUTMIX:\n                accuracy = (F.softmax(output).argmax(dim=1) == labels).float().mean()\n\n        # Uncomment below if using cutmix_ and fmix\n#         else:\n#             if labels[2] >= 0.5:\n#                 accuracy = (F.softmax(output).argmax(dim=1) == labels[0]).float().mean()\n#             else:\n#                 accuracy = (F.softmax(output).argmax(dim=1) == labels[1]).float().mean()\n        else:\n            if labels[0,2] >= 0.5:\n                accuracy = (output.argmax(dim=1) == labels[:,0]).float().mean()\n            else:\n                accuracy = (output.argmax(dim=1) == labels[:,1]).float().mean()\n\n        accuracies.update(accuracy.item(), BATCH_SIZE)\n        \n        if (step+1) % ITERS_TO_ACCUMULATE == 0:\n            # Unscale the gradients of the optimizer's assigned params. If these\n            # gradients do not contain infs or NaNs, optimizer.step() is then called,\n            # otherwise, optimizer.step() is skipped.\n            scaler.step(optimizer)\n            # Update the scale for next iteration.\n            scaler.update()\n            optimizer.zero_grad()\n        \n        if scheduler is not None and SCHEDULER_UPDATE == 'batch':\n            scheduler.step()\n\n        batch_time.update(time.time() - start_time)\n        start_time = time.time()\n        \n        if step % ITER_FREQ == 0:\n            \n            print('Epoch: [{0}][{1}/{2}]\\t'\n                  'Batch Time {batch_time.val:.3f}s ({batch_time.avg:.3f}s), '\n                  'Data Time {data_time.val:.3f}s ({data_time.avg:.3f}s)\\t'\n                  'Loss: {loss.val:.4f} ({loss.avg:.4f}), '\n                  'Accuracy {accuracy.val:.4f} ({accuracy.avg:.4f})'.format((epoch+1),\n                                                                             step, len(dataloader),\n                                                                             batch_time=batch_time,\n                                                                             data_time=data_time,\n                                                                             loss=losses,\n                                                                             accuracy=accuracies))\n        # To check the loss real-time while iterating over data.\n        loader.set_postfix(loss=losses.avg, accuracy=accuracies.avg)\n#         del images, labels\n    if scheduler is not None and SCHEDULER_UPDATE == 'epoch':\n        scheduler.step()\n        \n    return losses.avg, accuracies.avg","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def valid_fn(epoch, model, criterion, val_loader, device, scheduler):\n    \n    model.eval()\n    losses = AverageMeter()\n    accuracies = AverageMeter()\n    \n    loader = tqdm(val_loader, total=len(val_loader))\n    with torch.no_grad():  # without torch.no_grad() will make the CUDA run OOM.\n        for step, (images, labels) in enumerate(loader):\n\n            images = images.to(device)\n            labels = labels.to(device)\n\n            output = model(images)\n\n            loss = criterion(output, labels)\n            if CUTMIX:\n                losses.update(loss, BATCH_SIZE)\n            else:\n                losses.update(loss.item(), BATCH_SIZE)\n            accuracy = (F.softmax(output).argmax(dim=1) == labels).float().mean()\n            accuracies.update(accuracy.item(), VAL_BATCH_SIZE)\n\n            loader.set_postfix(loss=losses.avg, accuracy=accuracies.avg)\n#             del images, labels\n    \n    if scheduler is not None:\n        scheduler.step()\n        \n    return losses.avg, accuracies.avg","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"[Back to CFG(Click here)](#cont)"},{"metadata":{},"cell_type":"markdown","source":"# Main"},{"metadata":{"trusted":true,"scrolled":true},"cell_type":"code","source":"def engine(device, folds, fold, model_path=None):\n    \n    trn_idx = folds[folds['fold'] != fold].index\n    val_idx = folds[folds['fold'] == fold].index\n\n    train_folds = folds.loc[trn_idx].reset_index(drop=True)\n    valid_folds = folds.loc[val_idx].reset_index(drop=True)\n    \n    train_data = CLD_train_DS(train_folds, transform=get_transform())\n    val_data = CLD_train_DS(valid_folds, transform=get_transform(train=False))\n\n    if CUTMIX:\n        print('Using CutMix')\n        collator = CutMixCollator(CM_ALPHA)                \n    \n    train_loader = DataLoader(train_data,\n                              batch_size=BATCH_SIZE, \n                              shuffle=True, \n                              num_workers=NUM_WORKERS,\n#                               collate_fn=collator, \n                              pin_memory=True, # enables faster data transfer to CUDA-enabled GPUs.\n                              drop_last=True)\n    val_loader = DataLoader(val_data, \n                              batch_size=VAL_BATCH_SIZE,\n                              num_workers=NUM_WORKERS,\n                              shuffle=False, \n                              pin_memory=True,\n                              drop_last=False)\n    \n\n    if model_path is not None:\n        model = torch.load(model_path)\n        START_EPOCH = int(model_path.split('_')[-1])\n    else:\n        model = CustomEffNet(MODEL_ARCH, 5, True)\n        START_EPOCH = 0\n    model.to(device)\n    \n    params = filter(lambda p: p.requires_grad, model.parameters())    \n    optimizer = GetOptimizer(OPTIMIZER, params)\n    if not CUTMIX:\n        criterion = GetCriterion(CRITERION).to(device)\n    else:\n        criterion = GetCriterion('CutMix', GetCriterion(CRITERION)).to(device)\n        \n    val_criterion = GetCriterion(CRITERION).to(device)\n    scheduler = GetScheduler(SCHEDULER, optimizer)\n    \n    loss = []\n    accuracy = []\n    for epoch in range(START_EPOCH, EPOCHS):\n        \n        epoch_start = time.time()        \n        avg_loss, avg_accuracy = train_fn(model, train_loader, device, epoch, optimizer, criterion, scheduler)\n\n        torch.cuda.empty_cache()\n        avg_val_loss, avg_val_acc = valid_fn(epoch, model, val_criterion, val_loader, device, scheduler)\n        epoch_end = time.time() - epoch_start\n        \n        print(f'Validation accuracy after epoch {epoch+1}: {avg_val_acc:.4f}')\n        loss.append(avg_loss)\n        accuracy.append(avg_accuracy)\n        \n        content = f'Fold {fold} Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f} avg_train_accuracy: {avg_accuracy:.4f} avg_val_loss: {avg_val_loss:.4f} avg_val_acc: {avg_val_acc:.4f} time: {epoch_end:.0f}s'\n        with open(f'GPU_{MODEL_ARCH}_{OPTIMIZER}_{CRITERION}.txt', 'a') as appender:\n            appender.write(content + '\\n')\n        \n        # Save the model to use it for inference.\n        torch.save(model.state_dict(), f'GPU_{MODEL_ARCH}_fold_{fold}_epoch_{(epoch+1)}.pth')\n        torch.save(model, f'GPU_{MODEL_ARCH}_fold_{fold}_epoch_{(epoch+1)}')\n        torch.cuda.empty_cache()\n    \n    return {'loss':loss, 'accuracy':accuracy}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"scrolled":false},"cell_type":"code","source":"if __name__ == '__main__':\n    \n    if MODEL_PATH is not None:\n        START_FOLD = int(MODEL_PATH.split('_')[-3])\n    \n    for fold in range(START_FOLD, N_FOLDS):\n        print(f'===== Fold {fold} Starting =====')\n        fold_start = time.time()\n        logs = engine(DEVICE, folds, fold, MODEL_PATH)\n        print(f'Time taken in fold {fold}: {time.time()-fold_start}')","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}