{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","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'\n\nfrom 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\n\nCFG = {\n    'fold_num': 5,\n    'seed': 719,\n    'model_arch': 'tf_efficientnet_b4_ns',\n    'img_size': 512,\n    'epochs': 12,\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}\n\ntrain = pd.read_csv(f'{DATA_DIR}/train.csv')\ntrain.head()\n\ntrain.label.value_counts()\n\nfrom torch.optim.optimizer import Optimizer, required\nimport math\n\nclass AdamP(Optimizer):\n    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,\n                 weight_decay=0, delta=0.1, wd_ratio=0.1, nesterov=False):\n        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay,\n                        delta=delta, wd_ratio=wd_ratio, nesterov=nesterov)\n        super(AdamP, self).__init__(params, defaults)\n\n    def _channel_view(self, x):\n        return x.view(x.size(0), -1)\n\n    def _layer_view(self, x):\n        return x.view(1, -1)\n\n    def _cosine_similarity(self, x, y, eps, view_func):\n        x = view_func(x)\n        y = view_func(y)\n\n        return F.cosine_similarity(x, y, dim=1, eps=eps).abs_()\n\n    def _projection(self, p, grad, perturb, delta, wd_ratio, eps):\n        wd = 1\n        expand_size = [-1] + [1] * (len(p.shape) - 1)\n        for view_func in [self._channel_view, self._layer_view]:\n\n            cosine_sim = self._cosine_similarity(grad, p.data, eps, view_func)\n\n            if cosine_sim.max() < delta / math.sqrt(view_func(p.data).size(1)):\n                p_n = p.data / view_func(p.data).norm(dim=1).view(expand_size).add_(eps)\n                perturb -= p_n * view_func(p_n * perturb).sum(dim=1).view(expand_size)\n                wd = wd_ratio\n\n                return perturb, wd\n\n        return perturb, wd\n\n    def step(self, closure=None):\n        loss = None\n        if closure is not None:\n            loss = closure()\n\n        for group in self.param_groups:\n            for p in group['params']:\n                if p.grad is None:\n                    continue\n\n                grad = p.grad.data\n                beta1, beta2 = group['betas']\n                nesterov = group['nesterov']\n\n                state = self.state[p]\n\n                # State initialization\n                if len(state) == 0:\n                    state['step'] = 0\n                    state['exp_avg'] = torch.zeros_like(p.data)\n                    state['exp_avg_sq'] = torch.zeros_like(p.data)\n\n                # Adam\n                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n\n                state['step'] += 1\n                bias_correction1 = 1 - beta1 ** state['step']\n                bias_correction2 = 1 - beta2 ** state['step']\n\n                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)\n                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)\n\n                denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])\n                step_size = group['lr'] / bias_correction1\n\n                if nesterov:\n                    perturb = (beta1 * exp_avg + (1 - beta1) * grad) / denom\n                else:\n                    perturb = exp_avg / denom\n\n                # Projection\n                wd_ratio = 1\n                if len(p.shape) > 1:\n                    perturb, wd_ratio = self._projection(p, grad, perturb, group['delta'], group['wd_ratio'], group['eps'])\n\n                # Weight decay\n                if group['weight_decay'] > 0:\n                    p.data.mul_(1 - group['lr'] * group['weight_decay'] * wd_ratio)\n\n                # Step\n                p.data.add_(perturb, alpha=-step_size)\n\n        return loss\n\n\n\ndef 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()\n\ndef 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\n\ndef 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)\n\nclass 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.fc.in_features\n        self.model.fc = nn.Linear(n_features, n_class)\n\n    def forward(self, x):\n        x = self.model(x)\n        return x\n\ndef 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        \nclass LabelSmoothingLoss(torch.nn.Module):\n    def __init__(self, smoothing: float = 0.1, reduction=\"mean\", weight=None):\n        super(LabelSmoothingLoss, self).__init__()\n        self.epsilon = smoothing\n        self.reduction = reduction\n        self.weight = weight\n        \n\n    def reduce_loss(self, loss):\n        if self.reduction == \"mean\":\n            return loss.mean()\n        elif self.reduction == \"sum\":\n            return loss.sum()\n        else:\n            return loss\n\n    def linear_combination(self, x, y):\n        return self.epsilon * x + (1 - self.epsilon) * y\n\n    def forward(self, preds, target):\n        if self.weight is not None:\n            self.weight = self.weight.to(preds.device)\n\n        if self.training:\n            n = preds.size(-1)\n            log_preds = F.log_softmax(preds, dim=-1)\n            loss = self.reduce_loss(-log_preds.sum(dim=-1))\n            nll = F.nll_loss(\n                log_preds, target, reduction=self.reduction, weight=self.weight\n            )\n            return self.linear_combination(loss / n, nll)\n        else:\n            return torch.nn.functional.cross_entropy(preds, target, weight=self.weight)\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()\n\nclass 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\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 = AdamP(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 = LabelSmoothingLoss().to(device)\n        loss_fn = LabelSmoothingLoss().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(),'{}_fold_{}_{}'.format(CFG['model_arch'], fold, epoch))\n\n        del model, optimizer, train_loader, val_loader, scheduler, scaler\n        torch.cuda.empty_cache()","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}