{"cells":[{"metadata":{},"cell_type":"markdown","source":"# I just refactor the following [notebook](https://www.kaggle.com/khyeh0719/pytorch-efficientnet-baseline-train-amp-aug). "},{"metadata":{"trusted":true},"cell_type":"code","source":"package_paths = [\n    '../input/pytorch-image-models/pytorch-image-models-master',\n    '../input/image-fmix/FMix-master'\n]\nimport sys\n\nfor pth in package_paths:\n    sys.path.append(pth)\n\nDATA_DIR = '../input/cassava-leaf-disease-classification'","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"from datetime import datetime\nimport os\nimport random\nimport time\nimport warnings\n\nimport pandas as pd\nimport numpy as np\nimport cv2\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nimport sklearn\nfrom sklearn import metrics\nfrom sklearn.model_selection import StratifiedKFold\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.cuda.amp import autocast, GradScaler \nfrom torchvision import transforms\nfrom torch.utils.data import Dataset,DataLoader\nfrom torch.utils.data.sampler import SequentialSampler\nimport timm\n\nfrom albumentations import (\n    HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90,\n    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,\n    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine, RandomResizedCrop,\n    IAASharpen, IAAEmboss, RandomBrightnessContrast, Flip, OneOf, Compose, Normalize, Cutout, CoarseDropout, ShiftScaleRotate, CenterCrop, Resize\n)\nfrom albumentations.pytorch import ToTensorV2\nfrom fmix import make_low_freq_image, binarise_mask","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"CFG = {\n    'fold_num': 5,\n    'seed': 719,\n    'model_arch': 'tf_efficientnet_b4_ns',\n    'img_size': 512,\n    'epochs': 10,\n    'train_bs': 16,\n    'valid_bs': 32,\n    'T_0': 10,\n    'lr': 1e-4,\n    'min_lr': 1e-6,\n    'weight_decay':1e-6,\n    'num_workers': 4,\n    'accum_iter': 2,\n    'verbose_step': 1,\n    'device': 'cuda:0'\n}","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"train = pd.read_csv(f'{DATA_DIR}/train.csv')\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train.label.value_counts()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> We could do stratified validation split in each fold to make each fold's train and validation set looks like the whole train set in target distributions."},{"metadata":{"trusted":true},"cell_type":"code","source":"submission = pd.read_csv(f'{DATA_DIR}/sample_submission.csv')\nsubmission.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Helper Functions"},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True\n    \ndef get_img(path):\n    im_bgr = cv2.imread(path)\n    im_rgb = im_bgr[:, :, ::-1]\n    return im_rgb\n\nimg = get_img(f'{DATA_DIR}/train_images/1000015157.jpg')\nplt.imshow(img)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"def rand_bbox(size, lam):\n    W = size[0]\n    H = size[1]\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\n\nclass CassavaDataset(Dataset):\n    def __init__(self, df, data_root, \n                 transforms=None, \n                 output_label=True, \n                 one_hot_label=False,\n                 do_fmix=False, \n                 fmix_params={\n                     'alpha': 1., \n                     'decay_power': 3., \n                     'shape': (CFG['img_size'], CFG['img_size']),\n                     'max_soft': True, \n                     'reformulate': False\n                 },\n                 do_cutmix=False,\n                 cutmix_params={'alpha': 1,}\n                ):\n        \n        super().__init__()\n        self.df = df.reset_index(drop=True).copy()\n        self.transforms = transforms\n        self.data_root = data_root\n        self.do_fmix = do_fmix\n        self.fmix_params = fmix_params\n        self.do_cutmix = do_cutmix\n        self.cutmix_params = cutmix_params\n        \n        self.output_label = output_label\n        self.one_hot_label = one_hot_label\n        \n        if output_label:\n            self.labels = self.df['label'].values\n            \n            if one_hot_label:\n                self.labels = np.eye(self.df['label'].max()+1)[self.labels]\n            \n    def __len__(self):\n        return self.df.shape[0]\n    \n    def __getitem__(self, index: int):\n        \n        if self.output_label:\n            target = self.labels[index]\n          \n        img  = get_img(f\"{self.data_root}/{self.df.loc[index]['image_id']}\")\n\n        if self.transforms:\n            img = self.transforms(image=img)['image']\n        \n        if self.do_fmix and np.random.uniform(0., 1., size=1)[0] > 0.5:\n            with torch.no_grad():\n                lam = np.clip(np.random.beta(self.fmix_params['alpha'], self.fmix_params['alpha']), 0.6, 0.7)\n                \n                # Make mask, get mean / std\n                mask = make_low_freq_image(self.fmix_params['decay_power'], self.fmix_params['shape'])\n                mask = binarise_mask(mask, lam, self.fmix_params['shape'], self.fmix_params['max_soft'])\n    \n                fmix_ix = np.random.choice(self.df.index, size=1)[0]\n                fmix_img  = get_img(f\"{self.data_root}/{self.df.iloc[fmix_ix]['image_id']}\")\n\n                if self.transforms:\n                    fmix_img = self.transforms(image=fmix_img)['image']\n\n                mask_torch = torch.from_numpy(mask)\n                \n                img = mask_torch*img+(1.-mask_torch)*fmix_img\n                rate = mask.sum() / (CFG['img_size']**2)\n                target = rate*target + (1.-rate)*self.labels[fmix_ix]\n        \n        if self.do_cutmix and np.random.uniform(0., 1., size=1)[0] > 0.5:\n            with torch.no_grad():\n                cmix_ix = np.random.choice(self.df.index, size=1)[0]\n                cmix_img  = get_img(f\"{self.data_root}/{self.df.iloc[cmix_ix]['image_id']}\")\n                if self.transforms:\n                    cmix_img = self.transforms(image=cmix_img)['image']\n                    \n                lam = np.clip(np.random.beta(self.cutmix_params['alpha'], self.cutmix_params['alpha']), 0.3, 0.4)\n                bbx1, bby1, bbx2, bby2 = rand_bbox((CFG['img_size'], CFG['img_size']), lam)\n\n                img[:, bbx1:bbx2, bby1:bby2] = cmix_img[:, bbx1:bbx2, bby1:bby2]\n\n                rate = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (CFG['img_size']**2))\n                target = rate*target + (1.-rate)*self.labels[cmix_ix]\n                            \n        # do label smoothing\n        if self.output_label:\n            return img, target\n        else:\n            return img","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Define Train\\Validation Image Augmentations"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transforms():\n    return Compose([\n            RandomResizedCrop(CFG['img_size'], CFG['img_size']),\n            Transpose(p=0.5),\n            HorizontalFlip(p=0.5),\n            VerticalFlip(p=0.5),\n            ShiftScaleRotate(p=0.5),\n            HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, 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            CoarseDropout(p=0.5),\n            Cutout(p=0.5),\n            ToTensorV2(p=1.0),\n        ], p=1.0)\n  \n        \ndef get_valid_transforms():\n    return Compose([\n            CenterCrop(CFG['img_size'], CFG['img_size'], p=1.0),\n            Resize(CFG['img_size'], CFG['img_size']),\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.0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassvaImgClassifier(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.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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training APIs"},{"metadata":{"trusted":true},"cell_type":"code","source":"def prepare_dataloader(df, trn_idx, val_idx, data_root=f'{DATA_DIR}/train_images/'):\n    \n    from catalyst.data.sampler import BalanceClassSampler\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                              output_label=True, one_hot_label=False, do_fmix=False, do_cutmix=False)\n    valid_ds = CassavaDataset(valid_, data_root, transforms=get_valid_transforms(),output_label=True)\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\n\ndef train_one_epoch(epoch, model, loss_fn, optimizer, train_loader, device, scheduler=None, schd_batch_update=False):\n    model.train()\n\n    t = time.time()\n    running_loss = None\n\n    pbar = tqdm(enumerate(train_loader), total=len(train_loader))\n    for step, (imgs, image_labels) in pbar:\n        imgs = imgs.to(device).float()\n        image_labels = image_labels.to(device).long()\n\n        with autocast():\n            image_preds = model(imgs)\n\n            loss = loss_fn(image_preds, image_labels)\n            \n            scaler.scale(loss).backward()\n\n            if running_loss is None:\n                running_loss = loss.item()\n            else:\n                running_loss = running_loss * .99 + loss.item() * .01\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 scheduler is not None and schd_batch_update:\n                    scheduler.step()\n\n            if ((step + 1) % CFG['verbose_step'] == 0) or ((step + 1) == len(train_loader)):\n                description = f'epoch {epoch} loss: {running_loss:.4f}'\n                \n                pbar.set_description(description)\n                \n    if scheduler is not None and not schd_batch_update:\n        scheduler.step()\n        \ndef valid_one_epoch(epoch, model, loss_fn, val_loader, device, scheduler=None, schd_loss_update=False):\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, (imgs, image_labels) in pbar:\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    print(f'validation multi-class accuracy = {(image_preds_all==image_targets_all).mean():.4f}')\n    \n    if scheduler is not None:\n        if schd_loss_update:\n            scheduler.step(loss_sum/sample_num)\n        else:\n            scheduler.step()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Main Loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nif __name__ == '__main__':\n     # for training only, need nightly build pytorch\n\n    seed_everything(CFG['seed'])\n    \n    stratifiedKFold = StratifiedKFold(n_splits=CFG['fold_num'], shuffle=True, random_state=CFG['seed'])\n    folds = stratifiedKFold.split(np.arange(train.shape[0]), train.label.values)\n    \n    for fold, (trn_idx, val_idx) in enumerate(folds):\n        # we'll train fold 0 first\n        if fold > 0:\n            break\n\n        print(f'Training with {fold} started')\n        print(len(trn_idx), len(val_idx))\n\n        train_loader, val_loader = prepare_dataloader(train, trn_idx, val_idx, \n                                                      data_root=f'{DATA_DIR}/train_images/')\n\n        device = torch.device(CFG['device'])\n        \n        model = CassvaImgClassifier(CFG['model_arch'], train.label.nunique(), pretrained=True).to(device)\n        scaler = GradScaler()\n        optimizer = torch.optim.Adam(model.parameters(), lr=CFG['lr'], weight_decay=CFG['weight_decay'])\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n            optimizer, T_0=CFG['T_0'], T_mult=1, eta_min=CFG['min_lr'], last_epoch=-1)\n\n        loss_tr = nn.CrossEntropyLoss().to(device)\n        loss_fn = nn.CrossEntropyLoss().to(device)\n\n        for epoch in range(CFG['epochs']):\n            train_one_epoch(epoch, model, loss_tr, optimizer, train_loader, device,\n                            scheduler=scheduler, schd_batch_update=False)\n\n            with torch.no_grad():\n                valid_one_epoch(epoch, model, loss_fn, val_loader, device,\n                                scheduler=None, schd_loss_update=False)\n\n            torch.save(model.state_dict(), f\"{CFG['model_arch']}_fold_{fold}_{epoch}\")\n\n        del model, optimizer, train_loader, val_loader, scheduler, scaler\n        torch.cuda.empty_cache()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Inferece part is here: https://www.kaggle.com/hinamimi/efficientnet-baseline-inference-tta"}],"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}