{"cells":[{"metadata":{},"cell_type":"markdown","source":"# About this\n\nI have tried to implement the [Sharpness-Aware Minimization optimizer](https://arxiv.org/abs/2010.01412) in this Notebook, according to [this](https://github.com/davda54/sam) PyTorch implementation.\n\nI hope I got everything right, if anyone has some hints on this, please comment :)\n\nThe rest of the notebook is just playing aroung with some models and parameters, it actually uses a n-fold cv and is heavily influences by [this](https://www.kaggle.com/yasufuminakama/cassava-resnext50-32x4d-starter-training) notebook. "},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!pip install timm","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session\n#TRAIN_PATH = '../input/cassava-leaf-disease-classification/train_images'\nTRAIN_PATH = \"../input/cassava-leaf-disease-merged/train\"\nTEST_PATH = '../input/cassava-leaf-disease-classification/test_images'\n\nOUTPUT_DIR = './'","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"import sys\nsys.path.append('../input/pytorch-image-models/pytorch-image-models-master')\n\nimport math\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts\nfrom torch.utils.data.sampler import *\n\n\nfrom sklearn.metrics import accuracy_score\n\nimport cv2\n\nimport time\nimport timm\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"#train = pd.read_csv('../input/cassava-leaf-disease-classification/train.csv')\ntrain = pd.read_csv('../input/cassava-leaf-disease-merged/merged.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CFG:\n    debug=False\n    num_workers=8\n    model_name='tf_efficientnet_b6_ns' #resnext101_32x8d #resnext50_32x4d #tf_efficientnet_l2_ns\n    size=368\n    batch_size=16\n    seed=2002\n    target_size=5\n    target_col='label'\n    n_fold=3\n    trn_fold=[n for n in range(n_fold)]\n    inference=False\n    train=True\n    T_0=10\n    min_lr=1e-6\n    epochs = 10\n    max_grad_norm=1000\n    gradient_accumulation_steps=1\n    print_freq = 100\n\nif CFG.debug:\n    CFG.epochs = 1\n    train = train.sample(n=1000, random_state=CFG.seed).reset_index(drop=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from sklearn.model_selection import StratifiedKFold\n\nfolds = train.copy()\nFold = StratifiedKFold(n_splits=CFG.n_fold, shuffle=True, random_state=CFG.seed)\nfor n, (train_index, val_index) in enumerate(Fold.split(folds, folds[CFG.target_col])):\n    folds.loc[val_index, 'fold'] = int(n)\nfolds['fold'] = folds['fold'].astype(int)\nprint(folds.groupby(['fold', CFG.target_col]).size())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from albumentations.pytorch import ToTensorV2\nfrom albumentations import ImageOnlyTransform\nfrom albumentations import (\n    Compose, OneOf, Normalize, Resize, RandomResizedCrop, RandomCrop, HorizontalFlip, VerticalFlip, \n    RandomBrightness, RandomContrast, RandomBrightnessContrast, Rotate, ShiftScaleRotate, Cutout, \n    IAAAdditiveGaussianNoise, Transpose, HueSaturationValue\n    )\n\ndef get_transforms(*, data_type):\n    if data_type == \"train\":\n        return Compose([\n            RandomResizedCrop(CFG.size, CFG.size),\n            Transpose(p=0.5),\n            HorizontalFlip(p=0.5),\n            VerticalFlip(p=0.5),\n            ShiftScaleRotate(p=0.5),\n            HueSaturationValue(\n                hue_shift_limit=0.2, \n                sat_shift_limit=0.2,\n                val_shift_limit=0.2, \n                p=0.5\n            ),\n            RandomBrightnessContrast(\n                brightness_limit=(-0.1,0.1), \n                contrast_limit=(-0.1, 0.1), \n                p=0.5\n            ),\n            Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225],\n            ),\n            ToTensorV2(),\n        ])\n    elif data_type == 'valid':\n        return Compose([\n            Resize(CFG.size, CFG.size),\n            Normalize(\n                mean=[0.485, 0.456, 0.406],\n                std=[0.229, 0.224, 0.225],\n            ),\n            ToTensorV2(),\n        ])\n    ","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"class TrainDataset(Dataset):\n    def __init__(self, df, transform=None):\n        self.df = df\n        self.file_names = df['image_id'].values\n        self.labels = df['label'].values\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        file_name = self.file_names[idx]\n        file_path = f'{TRAIN_PATH}/{file_name}'\n        image = cv2.imread(file_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented['image']\n        label = torch.tensor(self.labels[idx]).long()\n        return image, label\n    \nclass TestDataset(Dataset):\n    def __init__(self, df, transform=None):\n        self.df = df\n        self.file_names = df['image_id'].values\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        file_name = self.file_names[idx]\n        file_path = f'{TEST_PATH}/{file_name}'\n        image = cv2.imread(file_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented['image']\n        return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CustomResNext(nn.Module):\n    def __init__(self, model_name=CFG.model_name, pretrained=False):\n        super().__init__()\n        self.model = timm.create_model(model_name, pretrained=pretrained)\n        n_features = self.model.fc.in_features\n        self.model.fc = nn.Linear(n_features, CFG.target_size)\n\n    def forward(self, x):\n        x = self.model(x)\n        return x\n\nclass CustomEffNet(nn.Module):\n    def __init__(self, model_name=CFG.model_name, pretrained=False):\n        super().__init__ ()\n        self.efnet_model = timm.create_model(model_name, pretrained=pretrained)\n        efnet_features = self.efnet_model.classifier.in_features\n        self.efnet_model.classifier = nn.Linear(efnet_features, CFG.target_size)\n\n        # self.resnet_model = timm.create_model(resnet_arch, pretrained=pretrained)\n        # resnet_features = self.resnet_model.fc.in_features\n        # self.resnet_model.classifier = nn.Linear(resnet_features, n_class)\n\n    def forward(self, x):\n        efnet_opt = self.efnet_model(x)\n        # resnet_opt = self.resnet_model(x)\n\n        # combined = efnet_opt*0.6 + resnet_opt*0.4\n        return efnet_opt\n\nclass EnsembleClassifier(nn.Module):\n    def __init__ (self, pretrained=False):\n        super().__init__()\n        self.model1 = timm.create_model(\"resnext50_32x4d\", pretrained=pretrained)\n        self.model2 = timm.create_model(\"tf_efficientnet_b7_ns\", pretrained=pretrained)\n        \n    def forward(self, x): \n        x1 = self.model1(x)\n        x2 = self.model2(x)\n        return x1 * 0.5 + x2 * 0.5\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Here comes the SAM algorithm :)"},{"metadata":{"trusted":true},"cell_type":"code","source":"base_optimizer = torch.optim.SGD\n\nclass 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    def step(self, closure=None):\n        raise NotImplementedError(\"SAM doesn't work like the other optimizers, you should first call `first_step` and the `second_step`; see the documentation for more info.\")\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":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef smooth_crossentropy(pred, gold, smoothing=0.1):\n    n_class = pred.size(1)\n\n    one_hot = torch.full_like(pred, fill_value=smoothing / (n_class - 1))\n    one_hot.scatter_(dim=1, index=gold.unsqueeze(1), value=1.0 - smoothing)\n    log_prob = F.log_softmax(pred, dim=1)\n\n    return F.kl_div(input=log_prob, target=one_hot, reduction='none').sum(-1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class AverageMeter:\n    \"\"\"\n    Computes and stores the average and current value\n    \"\"\"\n\n    def __init__(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 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 asMinutes(s):\n    m = math.floor(s / 60)\n    s -= m * 60\n    return '%dm %ds' % (m, s)\n\n\ndef timeSince(since, percent):\n    now = time.time()\n    s = now - since\n    es = s / (percent)\n    rs = es - s\n    return '%s (remain %s)' % (asMinutes(s), asMinutes(rs))\n\ndef init_logger(log_file=OUTPUT_DIR+'train.log'):\n    from logging import getLogger, INFO, FileHandler,  Formatter,  StreamHandler\n    logger = getLogger(__name__)\n    logger.setLevel(INFO)\n    handler1 = StreamHandler()\n    handler1.setFormatter(Formatter(\"%(message)s\"))\n    handler2 = FileHandler(filename=log_file)\n    handler2.setFormatter(Formatter(\"%(message)s\"))\n    logger.addHandler(handler1)\n    logger.addHandler(handler2)\n    return logger\n\nLOGGER = init_logger()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Train Function**"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_fn(train_loader, model, criterion, optimizer, epoch, scheduler, device):\n    batch_time = AverageMeter()\n    data_time = AverageMeter()\n    losses = AverageMeter()\n    scores = AverageMeter()\n    # switch to train mode\n    model.train()\n    start = end = time.time()\n    global_step = 0\n    for step, (images, labels) in enumerate(train_loader):\n        # measure data loading time\n        data_time.update(time.time() - end)\n        images = images.to(device)\n        labels = labels.to(device)\n        batch_size = labels.size(0)\n        \n        # first forward-backward pass\n        with torch.cuda.amp.autocast():\n            preds_first = model(images)\n            loss = criterion(preds_first, labels)  # use this loss for any training statistics\n        \n        loss.mean().backward()\n        optimizer.first_step(zero_grad=True)\n        \n        # second forward-backward pass\n        with torch.cuda.amp.autocast():\n            preds_second = model(images)\n            loss_second = criterion(preds_second, labels)\n            \n        loss_second.mean().backward()\n        optimizer.second_step(zero_grad=True)\n        \n        #print(loss)\n        \n        losses.update(loss.item(), batch_size)\n        \n        \"\"\"\n        #loss = criterion(y_preds, labels)\n        # record loss\n        \n        if CFG.gradient_accumulation_steps > 1:\n            loss = loss / CFG.gradient_accumulation_steps\n        if CFG.apex:\n            with amp.scale_loss(loss, optimizer) as scaled_loss:\n                scaled_loss.backward()\n        else:\n            loss.backward()\n        \n        \n        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), CFG.max_grad_norm)\n        \n        \n        if (step + 1) % CFG.gradient_accumulation_steps == 0:\n            optimizer.step()\n            optimizer.zero_grad()\n            global_step += 1\n        \"\"\"\n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n        if step % CFG.print_freq == 0 or step == (len(train_loader)-1):\n            print('Epoch: [{0}][{1}/{2}] '\n                  'Data {data_time.val:.3f} ({data_time.avg:.3f}) '\n                  'Elapsed {remain:s} '\n                  'Loss: {loss.val:.4f}({loss.avg:.4f}) '\n                  #'Grad: {grad_norm:.4f}  '\n                  'LR: {lr:.6f}  '\n                  .format(\n                   epoch+1, step, len(train_loader), batch_time=batch_time,\n                   data_time=data_time, loss=losses,\n                   remain=timeSince(start, float(step+1)/len(train_loader)),\n                   #grad_norm=grad_norm,\n                   lr=scheduler.get_last_lr()[0],\n                   ))\n    return losses.avg","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def valid_fn(valid_loader, model, criterion, device):\n    batch_time = AverageMeter()\n    data_time = AverageMeter()\n    losses = AverageMeter()\n    scores = AverageMeter()\n    # switch to evaluation mode\n    model.eval()\n    preds = []\n    start = end = time.time()\n    for step, (images, labels) in enumerate(valid_loader):\n        # measure data loading time\n        data_time.update(time.time() - end)\n        images = images.to(device)\n        labels = labels.to(device)\n        batch_size = labels.size(0)\n        # compute loss\n        with torch.no_grad():\n            y_preds = model(images)\n        loss = criterion(y_preds, labels)\n        losses.update(loss.item(), batch_size)\n        # record accuracy\n        preds.append(y_preds.softmax(1).to('cpu').numpy())\n        if CFG.gradient_accumulation_steps > 1:\n            loss = loss / CFG.gradient_accumulation_steps\n        # measure elapsed time\n        batch_time.update(time.time() - end)\n        end = time.time()\n        if step % CFG.print_freq == 0 or step == (len(valid_loader)-1):\n            print('EVAL: [{0}/{1}] '\n                  'Data {data_time.val:.3f} ({data_time.avg:.3f}) '\n                  'Elapsed {remain:s} '\n                  'Loss: {loss.val:.4f}({loss.avg:.4f}) '\n                  .format(\n                   step, len(valid_loader), batch_time=batch_time,\n                   data_time=data_time, loss=losses,\n                   remain=timeSince(start, float(step+1)/len(valid_loader)),\n                   ))\n    predictions = np.concatenate(preds)\n    return losses.avg, predictions","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from tqdm import tqdm\n\ndef inference(model, states, test_loader, device):\n    model.to(device)\n    tk0 = tqdm(enumerate(test_loader), total=len(test_loader))\n    probs = []\n    for i, (images) in tk0:\n        images = images.to(device)\n        avg_preds = []\n        for state in states:\n            model.load_state_dict(state['model'])\n            model.eval()\n            with torch.no_grad():\n                y_preds = model(images)\n            avg_preds.append(y_preds.softmax(1).to('cpu').numpy())\n        avg_preds = np.mean(avg_preds, axis=0)\n        probs.append(avg_preds)\n    probs = np.concatenate(probs)\n    return probs","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_loop(folds, fold):\n    # ====================================================\n    # loader\n    # ====================================================\n    LOGGER.info(f\"========== Training fold: {fold+1}  ==========\")\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_dataset = TrainDataset(train_folds, \n                                 transform=get_transforms(data_type='train'))\n    valid_dataset = TrainDataset(valid_folds, \n                                 transform=get_transforms(data_type='valid'))\n\n    train_loader = DataLoader(train_dataset, \n                              batch_size=CFG.batch_size, \n                              shuffle=True, \n                              num_workers=CFG.num_workers, pin_memory=True, drop_last=True)\n    \n    valid_loader = DataLoader(valid_dataset, \n                              batch_size=CFG.batch_size, \n                              shuffle=False, \n                              num_workers=CFG.num_workers, pin_memory=True, drop_last=False)\n    \n    \n    #model = CustomResNext(CFG.model_name, pretrained=True)\n    model = CustomEffNet(CFG.model_name, pretrained = True)\n    model.to(device)\n    \n    #optimizer = Adam(model.parameters(), lr=CFG.lr, weight_decay=CFG.weight_decay, amsgrad=False)\n    optimizer = SAM(model.parameters(), base_optimizer, lr=0.1, momentum=0.9,weight_decay=0.0005)\n    \n    scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=CFG.T_0, T_mult=1, eta_min=CFG.min_lr, last_epoch=-1)\n    \n    loss_train = nn.CrossEntropyLoss()\n    loss_valid = nn.CrossEntropyLoss()\n    \n    best_score = 0.\n    best_loss = np.inf\n\n    \n    for epoch in range(CFG.epochs):\n        start_time = time.time()\n        \n        avg_loss = train_fn(train_loader, model, loss_train, optimizer, epoch, scheduler, device)\n        avg_val_loss, preds = valid_fn(valid_loader, model, loss_valid, device)\n        valid_labels = valid_folds[CFG.target_col].values\n        \n        scheduler.step()\n        \n        score = accuracy_score(valid_labels, preds.argmax(1))\n        \n        elapsed = time.time() - start_time\n        \n        LOGGER.info(f'Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  avg_val_loss: {avg_val_loss:.4f}  time: {elapsed:.0f}s')\n        LOGGER.info(f'Epoch {epoch+1} - Accuracy: {score}')\n        \n        if score > best_score:\n            best_score = score\n            LOGGER.info(f'Epoch {epoch+1} - Save Best Score: {best_score:.4f} Model')\n            torch.save({'model': model.state_dict(), \n                        'preds': preds},\n                        OUTPUT_DIR+f'{CFG.model_name}_fold{fold}_best.pth')\n            \n    check_point = torch.load(OUTPUT_DIR+f'{CFG.model_name}_fold{fold}_best.pth')\n    valid_folds[[str(c) for c in range(5)]] = check_point['preds']\n    valid_folds['preds'] = check_point['preds'].argmax(1)\n    \n    return valid_folds ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def main():\n\n    \"\"\"\n    Prepare: 1.train  2.test  3.submission  4.folds\n    \"\"\"\n\n    def get_result(result_df):\n        preds = result_df['preds'].values\n        labels = result_df[CFG.target_col].values\n        score = accuracy_score(labels, preds)\n        LOGGER.info(f'Score: {score:<.5f}')\n    \n    if CFG.train:\n        # train \n        oof_df = pd.DataFrame()\n        for fold in range(CFG.n_fold):\n            if fold in CFG.trn_fold:\n                _oof_df = train_loop(folds, fold)\n                oof_df = pd.concat([oof_df, _oof_df])\n                LOGGER.info(f\"========== fold: {fold} result ==========\")\n                get_result(_oof_df)\n        # CV result\n        LOGGER.info(f\"========== CV ==========\")\n        get_result(oof_df)\n        # save result\n        oof_df.to_csv(OUTPUT_DIR+'oof_df.csv', index=False)\n    \n    if CFG.inference:\n        # inference\n        #model = CustomResNext(CFG.model_name, pretrained=False) #Custom ResNext\n        \n        model =  CustomEffNet(CFG.model_name, pretrained = False) #EffNet\n        states = [torch.load(OUTPUT_DIR+f'{CFG.model_name}_fold{fold}_best.pth') for fold in CFG.trn_fold]\n        test_dataset = TestDataset(test, transform=get_transforms(data_type='valid'))\n        test_loader = DataLoader(test_dataset, batch_size=CFG.batch_size, shuffle=False, \n                                 num_workers=CFG.num_workers, pin_memory=True)\n        predictions = inference(model, states, test_loader, device)\n        # submission\n        test['label'] = predictions.argmax(1)\n        print(test)\n        test[['image_id', 'label']].to_csv(OUTPUT_DIR+'submission.csv', index=False)\n        \nif __name__ == '__main__':\n    main()","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}