{"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_minor":4,"nbformat":4,"cells":[{"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":{}},{"cell_type":"code","source":"!pip install timm","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:49.957409Z","iopub.execute_input":"2023-04-05T20:15:49.957750Z","iopub.status.idle":"2023-04-05T20:15:56.367559Z","shell.execute_reply.started":"2023-04-05T20:15:49.957717Z","shell.execute_reply":"2023-04-05T20:15:56.366259Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","execution":{"iopub.status.busy":"2023-04-05T20:15:56.374713Z","iopub.execute_input":"2023-04-05T20:15:56.375823Z","iopub.status.idle":"2023-04-05T20:15:56.387796Z","shell.execute_reply.started":"2023-04-05T20:15:56.375780Z","shell.execute_reply":"2023-04-05T20:15:56.386862Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from albumentations import Compose\nfrom albumentations.pytorch import ToTensorV2","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.390719Z","iopub.execute_input":"2023-04-05T20:15:56.391641Z","iopub.status.idle":"2023-04-05T20:15:56.399959Z","shell.execute_reply.started":"2023-04-05T20:15:56.391604Z","shell.execute_reply":"2023-04-05T20:15:56.399079Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Ignore warnings\nimport warnings\nwarnings.filterwarnings(\"ignore\")","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.404307Z","iopub.execute_input":"2023-04-05T20:15:56.404815Z","iopub.status.idle":"2023-04-05T20:15:56.416131Z","shell.execute_reply.started":"2023-04-05T20:15:56.404778Z","shell.execute_reply":"2023-04-05T20:15:56.414951Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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'","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.418512Z","iopub.execute_input":"2023-04-05T20:15:56.418937Z","iopub.status.idle":"2023-04-05T20:15:56.428273Z","shell.execute_reply.started":"2023-04-05T20:15:56.418889Z","shell.execute_reply":"2023-04-05T20:15:56.427070Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# TODO : Play around with SNAPMIX_PCT, SNAPMIX_ALPHA and EPOCHS to converge for max accuracy","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.430402Z","iopub.execute_input":"2023-04-05T20:15:56.432155Z","iopub.status.idle":"2023-04-05T20:15:56.448112Z","shell.execute_reply.started":"2023-04-05T20:15:56.432111Z","shell.execute_reply":"2023-04-05T20:15:56.447136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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)    ","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.450702Z","iopub.execute_input":"2023-04-05T20:15:56.451229Z","iopub.status.idle":"2023-04-05T20:15:56.461043Z","shell.execute_reply.started":"2023-04-05T20:15:56.451196Z","shell.execute_reply":"2023-04-05T20:15:56.460205Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Cassava Dataset","metadata":{}},{"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","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.463099Z","iopub.execute_input":"2023-04-05T20:15:56.464599Z","iopub.status.idle":"2023-04-05T20:15:56.488130Z","shell.execute_reply.started":"2023-04-05T20:15:56.464561Z","shell.execute_reply":"2023-04-05T20:15:56.487268Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df = pd.read_csv(DATA_PATH + \"train.csv\")","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.489554Z","iopub.execute_input":"2023-04-05T20:15:56.490200Z","iopub.status.idle":"2023-04-05T20:15:56.550278Z","shell.execute_reply.started":"2023-04-05T20:15:56.490162Z","shell.execute_reply":"2023-04-05T20:15:56.549019Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Transforms using albumentations","metadata":{}},{"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            A.RandomBrightness(limit=0.3, p=0.5),\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            A.RandomBrightness(limit=0.3, p=0.5),\n            ToTensorV2(p=1.0),\n        ], p=1.)","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.551803Z","iopub.execute_input":"2023-04-05T20:15:56.552369Z","iopub.status.idle":"2023-04-05T20:15:56.573620Z","shell.execute_reply.started":"2023-04-05T20:15:56.552328Z","shell.execute_reply":"2023-04-05T20:15:56.572642Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model (modified to support SnapMix)","metadata":{}},{"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","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.576509Z","iopub.execute_input":"2023-04-05T20:15:56.577595Z","iopub.status.idle":"2023-04-05T20:15:56.602191Z","shell.execute_reply.started":"2023-04-05T20:15:56.577552Z","shell.execute_reply":"2023-04-05T20:15:56.601065Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Metrics","metadata":{}},{"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","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.603748Z","iopub.execute_input":"2023-04-05T20:15:56.605841Z","iopub.status.idle":"2023-04-05T20:15:56.623665Z","shell.execute_reply.started":"2023-04-05T20:15:56.605797Z","shell.execute_reply":"2023-04-05T20:15:56.622637Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from sklearn.metrics import precision_recall_fscore_support\n\ndef prec_recall_fscore(input, targs):\n    scores = precision_recall_fscore_support(targs.cpu(), input.cpu(), average='macro')\n    return {'precision': scores[0], 'recall': scores[1], 'fscore': scores[2], 'support': scores[3]} # return dictionary of 4 values\n\ndef print_prec_recall_fscore(scores):\n    print(\"Precision Recall FScores Metric:\")\n    print(scores)\n    \n    return scores","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.625226Z","iopub.execute_input":"2023-04-05T20:15:56.625798Z","iopub.status.idle":"2023-04-05T20:15:56.639959Z","shell.execute_reply.started":"2023-04-05T20:15:56.625758Z","shell.execute_reply":"2023-04-05T20:15:56.638879Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.641404Z","iopub.execute_input":"2023-04-05T20:15:56.642091Z","iopub.status.idle":"2023-04-05T20:15:56.649687Z","shell.execute_reply.started":"2023-04-05T20:15:56.642050Z","shell.execute_reply":"2023-04-05T20:15:56.648752Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Checkpoint method","metadata":{}},{"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))","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.652682Z","iopub.execute_input":"2023-04-05T20:15:56.654118Z","iopub.status.idle":"2023-04-05T20:15:56.667925Z","shell.execute_reply.started":"2023-04-05T20:15:56.654059Z","shell.execute_reply":"2023-04-05T20:15:56.666785Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Create folds","metadata":{}},{"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)","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.670619Z","iopub.execute_input":"2023-04-05T20:15:56.672144Z","iopub.status.idle":"2023-04-05T20:15:56.683536Z","shell.execute_reply.started":"2023-04-05T20:15:56.672006Z","shell.execute_reply":"2023-04-05T20:15:56.682667Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# SnapMix Augmentation","metadata":{}},{"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()","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.685014Z","iopub.execute_input":"2023-04-05T20:15:56.685630Z","iopub.status.idle":"2023-04-05T20:15:56.786559Z","shell.execute_reply.started":"2023-04-05T20:15:56.685592Z","shell.execute_reply":"2023-04-05T20:15:56.785489Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# SnapMix Criterion (Loss)","metadata":{}},{"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","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.788533Z","iopub.execute_input":"2023-04-05T20:15:56.792053Z","iopub.status.idle":"2023-04-05T20:15:56.801948Z","shell.execute_reply.started":"2023-04-05T20:15:56.792014Z","shell.execute_reply":"2023-04-05T20:15:56.801013Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# ADDED: FOCAL LOSS","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.autograd import Variable\n\nclass FocalLoss(nn.Module):\n    def __init__(self, gamma=0, alpha=None, size_average=True):\n        super(FocalLoss, self).__init__()\n        self.gamma = gamma\n        self.alpha = alpha\n        if isinstance(alpha,(float,int,int)): self.alpha = torch.Tensor([alpha,1-alpha])\n        if isinstance(alpha,list): self.alpha = torch.Tensor(alpha)\n        self.size_average = size_average\n\n    def forward(self, input, target):\n        if input.dim()>2:\n            input = input.view(input.size(0),input.size(1),-1)  # N,C,H,W => N,C,H*W\n            input = input.transpose(1,2)    # N,C,H*W => N,H*W,C\n            input = input.contiguous().view(-1,input.size(2))   # N,H*W,C => N*H*W,C\n        target = target.view(-1,1)\n\n        logpt = F.log_softmax(input)\n        logpt = logpt.gather(1,target)\n        logpt = logpt.view(-1)\n        pt = Variable(logpt.data.exp())\n\n        if self.alpha is not None:\n            if self.alpha.type()!=input.data.type():\n                self.alpha = self.alpha.type_as(input.data)\n            at = self.alpha.gather(0,target.data.view(-1))\n            logpt = logpt * Variable(at)\n\n        loss = -1 * (1-pt)**self.gamma * logpt\n        if self.size_average: return loss.mean()\n        else: return loss.sum()","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.804757Z","iopub.execute_input":"2023-04-05T20:15:56.806063Z","iopub.status.idle":"2023-04-05T20:15:56.847045Z","shell.execute_reply.started":"2023-04-05T20:15:56.806026Z","shell.execute_reply":"2023-04-05T20:15:56.846117Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train & Validate","metadata":{}},{"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    # criterion = FocalLoss().to(device)\n    # val_criterion = FocalLoss().to(device)\n    snapmix_criterion = SnapMixLoss().to(device)\n    \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                    # loss = torch.mean(criterion(outputs, y))                        \n                else:\n                    outputs, _ = model(X)\n                    loss = torch.mean(criterion(outputs, y))                \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['FLtrain_loss'] = round(train_loss/cur_step, 4)\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        # print_prec_recall_fscore(prec_recall_scores)\n        \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","metadata":{"execution":{"iopub.status.busy":"2023-04-05T20:15:56.848888Z","iopub.execute_input":"2023-04-05T20:15:56.849664Z","iopub.status.idle":"2023-04-05T21:35:23.444681Z","shell.execute_reply.started":"2023-04-05T20:15:56.849628Z","shell.execute_reply":"2023-04-05T21:35:23.442057Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print_prec_recall_fscore(prec_recall_fscore(preds, y))","metadata":{"execution":{"iopub.status.busy":"2023-04-05T21:35:23.446965Z","iopub.execute_input":"2023-04-05T21:35:23.447699Z","iopub.status.idle":"2023-04-05T21:35:23.459587Z","shell.execute_reply.started":"2023-04-05T21:35:23.447655Z","shell.execute_reply":"2023-04-05T21:35:23.458669Z"},"trusted":true},"execution_count":null,"outputs":[]}]}