{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Training pipline with ViT using Pytorch \nThis is a pipeline on training with ViT using PyTorch. If anyone finds any improvement, please comment in the notebook!\n\nReferences:\n\n[paper](https://arxiv.org/abs/2010.11929)\n\n[Github](https://github.com/rwightman/pytorch-image-models)"},{"metadata":{},"cell_type":"markdown","source":"# Install Timm"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"!pip -q install ../input/timm031py3noneanywhl/timm-0.3.1-py3-none-any.whl","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Print GPU info"},{"metadata":{"trusted":true},"cell_type":"code","source":"gpu_info = !nvidia-smi\ngpu_info = '\\n'.join(gpu_info)\nif gpu_info.find('failed') >= 0:\n    print('Select the Runtime > \"Change runtime type\" menu to enable a GPU accelerator, ')\n    print('and then re-execute this cell.')\nelse:\n    print(gpu_info)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **Import 3rdparty**"},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\nimport random\nimport cv2\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader, Dataset\nimport torch.nn.functional as F\nfrom torch.cuda.amp import autocast, GradScaler\n\nimport timm\nfrom matplotlib import pyplot as plt\nfrom sklearn.model_selection import StratifiedKFold\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom tqdm import tqdm","execution_count":null,"outputs":[]},{"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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **Global config**"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Config:\n    seed = 42\n    data_dir = '../input/cassava-leaf-disease-classification/'\n    train_data_dir = data_dir + 'train_images/'\n    train_csv_path = data_dir + 'train.csv'\n    arch = 'vit_base_patch16_384' ## model name\n    device = 'cuda'\n    debug = True                 ##\n    \n    image_size = 384     \n    train_batch_size = 16\n    val_batch_size = 32\n    epochs = 10                 ## total train epochs\n    freeze_bn_epochs = 5        ## freeze bn weights before epochs\n    \n    lr=1e-4                     ## init learning rate\n    min_lr = 1e-6               ## min learning rate\n    weight_decay = 1e-6\n    num_workers = 4\n    num_splits = 5             ## numbers splits\n    num_classes = 5            ## numbers classes\n    T_0 = 10\n    T_mult = 1\n    accum_iter = 2\n    verbose_step = 1\n    \n    criterion = 'LabelSmoothingCrossEntropy' ## CrossEntropy, LabelSmoothingCrossEntropy\n    label_smoothing = 0.3\n    \n    train_id = [0,1,2,3,4]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **Load Image**"},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_image(image_path):\n    img = cv2.imread(image_path)\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    return img","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **CassavaDataset**"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaDataset(Dataset):\n    def __init__(self, data_dir, df, transforms=None, output_label=True):\n        self.data_dir = data_dir\n        self.df = df\n        self.transforms = transforms\n        self.output_label = output_label\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, index):\n        image_infos = self.df.iloc[index]\n        image_path = self.data_dir + image_infos.image_id\n\n        image = load_image(image_path)\n\n        if image is None:\n            raise FileNotFoundError(image_path)\n\n        ### augment\n        if self.transforms is not None:\n            image = self.transforms(image=image)['image']\n        else:\n            image = torch.from_numpy(image)\n\n        if self.output_label:\n            return image, image_infos.label\n        else:\n            return image","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **CassavaClassifier**"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaClassifier(nn.Module):\n    def __init__(self, model_arch, num_classes, pretrained=False):\n        super().__init__()\n        self.model = timm.create_model(model_arch, pretrained=pretrained)\n        ### vit\n        num_features = self.model.head.in_features\n        self.model.head = nn.Linear(num_features, num_classes)\n\n        \n        '''\n        self.model.classifier = nn.Sequential(\n            nn.Dropout(0.3),\n            #nn.Linear(num_features, hidden_size,bias=True), nn.ELU(),\n            nn.Linear(num_features, num_classes, bias=True)\n        )\n        '''\n    def forward(self, x):\n        x = self.model(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **Train and Val transforms**"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_train_transforms(CFG):\n    return A.Compose([\n            A.RandomResizedCrop(height=CFG.image_size, width=CFG.image_size, p=0.5),\n            A.Transpose(p=0.5),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            A.RandomRotate90(p=0.5),\n            A.ShiftScaleRotate(p=0.5),\n            A.HueSaturationValue(hue_shift_limit=0.2, sat_shift_limit=0.2, val_shift_limit=0.2, p=0.5),\n            A.RandomBrightnessContrast(brightness_limit=(-0.1,0.1), contrast_limit=(-0.1, 0.1), p=0.5),\n            A.CenterCrop(CFG.image_size, CFG.image_size),\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.CoarseDropout(p=0.5),\n            A.Cutout(p=0.5),\n            ToTensorV2(),\n        ],p=1.0)\n\ndef get_val_transforms(cfg):\n    return A.Compose([\n            A.CenterCrop(CFG.image_size, CFG.image_size, p=0.5),\n            A.Resize(CFG.image_size, CFG.image_size),\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            ToTensorV2(),\n        ],p=1.0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **Train and Val data loader**"},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_dataloader(CFG, df, train_idx, val_idx):\n    df_train = df.loc[train_idx,:].reset_index(drop=True)\n    df_val = df.loc[val_idx,:].reset_index(drop=True)\n\n    train_dataset = CassavaDataset(\n        CFG.train_data_dir,\n        df_train,\n        transforms=get_train_transforms(CFG), \n        output_label=True)\n\n    val_dataset = CassavaDataset(\n        CFG.train_data_dir,\n        df_val,\n        transforms=get_val_transforms(CFG), \n        output_label=True)\n\n    train_loader = torch.utils.data.DataLoader(\n        train_dataset,\n        batch_size=CFG.train_batch_size,\n        pin_memory=False,\n        drop_last=False,\n        shuffle=True,        \n        num_workers=CFG.num_workers,\n        #sampler=BalanceClassSampler(labels=train_['label'].values, mode=\"downsampling\")\n    )\n\n    val_loader = torch.utils.data.DataLoader(\n        val_dataset, \n        batch_size=CFG.val_batch_size,\n        num_workers=CFG.num_workers,\n        shuffle=False,\n        pin_memory=False,\n    )\n    \n    return train_loader, val_loader","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **Train one epoch**"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_one_epoch(epoch,model,loss_fn,optimizer,train_loader,device,scheduler=None,schd_batch_update=False):\n    model.train()\n    lr = optimizer.state_dict()['param_groups'][0]['lr']\n    \n    running_loss = None\n    pbar = tqdm(enumerate(train_loader),total=len(train_loader))\n    for step,(images,targets) in pbar:\n        images = images.to(device).float()\n        targets = targets.to(device).long()\n        \n        with autocast():\n            preds = model(images)\n            loss = loss_fn(preds,targets)\n        \n            scaler.scale(loss).backward()\n            if running_loss is None:\n                running_loss = loss.item()\n            else:\n                running_loss = running_loss* 0.99 + loss.item()*0.01\n                \n            if ((step + 1) % CFG.accum_iter == 0) or ((step + 1) == len(train_loader)):\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            if ((step + 1) % CFG.accum_iter == 0) or ((step + 1) == len(train_loader)):\n                description = f'Train epoch {epoch} loss: {running_loss:.5f}'\n                pbar.set_description(description)\n                \n    if scheduler is not None and schd_batch_update:\n        scheduler.step()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **Valid one epoch**"},{"metadata":{"trusted":true},"cell_type":"code","source":"def valid_one_epoch(epoch,model,loss_fn,val_loader,device,scheduler=None,schd_loss_update=False):\n    model.eval()\n    \n    loss_sum = 0\n    sample_num = 0\n    preds_all = []\n    targets_all = []\n    scores = []\n    \n    pbar = tqdm(enumerate(val_loader),total=len(val_loader))\n    for step,(images,targets) in pbar:\n        images = images.to(device).float()\n        targets = targets.to(device).long()\n        preds = model(images)\n            \n        preds_all += [torch.argmax(preds,1).detach().cpu().numpy()]\n        targets_all += [targets.detach().cpu().numpy()]\n\n        loss = loss_fn(preds,targets)\n        loss_sum += loss.item()*targets.shape[0]\n        sample_num += targets.shape[0]\n           \n        if ((step + 1) % CFG.accum_iter == 0) or ((step + 1) == len(train_loader)):\n            description = f'Val epoch {epoch} loss: {loss_sum/sample_num:.5f}'\n            pbar.set_description(description)\n            \n    preds_all = np.concatenate(preds_all)\n    targets_all = np.concatenate(targets_all)\n    accuracy = (preds_all == targets_all).mean()\n    print(f'Validation multi-class accuracy = {accuracy:.5f}')\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    \n    return accuracy","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **Freeze bn weights**"},{"metadata":{"trusted":true},"cell_type":"code","source":"################ freeze bn \ndef freeze_batchnorm_stats(net):\n    try:\n        for m in net.modules():\n            if isinstance(m,nn.BatchNorm2d) or isinstance(m,nn.LayerNorm):\n                m.eval()\n    except ValuError:\n        print('error with batchnorm2d or layernorm')\n        return","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **Label Smoothing Cross Entropy Loss**"},{"metadata":{"trusted":true},"cell_type":"code","source":"class LabelSmoothingCrossEntropy(nn.Module):\n    \"\"\"\n    NLL loss with label smoothing.\n    \"\"\"\n    def __init__(self, smoothing=0.1):\n        \"\"\"\n        Constructor for the LabelSmoothing module.\n        :param smoothing: label smoothing factor\n        \"\"\"\n        super(LabelSmoothingCrossEntropy, self).__init__()\n        assert smoothing < 1.0\n        self.smoothing = smoothing\n        self.confidence = 1. - smoothing\n\n    def forward(self, x, target):\n        logprobs = F.log_softmax(x, dim=-1)\n        nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1))\n        nll_loss = nll_loss.squeeze(1)\n        smooth_loss = -logprobs.mean(dim=-1)\n        loss = self.confidence * nll_loss + self.smoothing * smooth_loss\n        return loss.mean()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# **Main Loop**"},{"metadata":{"trusted":true},"cell_type":"code","source":"if __name__ == '__main__':\n    CFG = Config\n    train = pd.read_csv(CFG.train_csv_path)\n    \n    if CFG.debug:\n        CFG.epochs = 1\n        train = train.sample(100,random_state=CFG.seed).reset_index(drop=True)\n    \n    print('CFG seed is ', CFG.seed)\n    if CFG.seed is not None:\n        seed_everything(CFG.seed)\n    \n    folds = StratifiedKFold(\n        n_splits=CFG.num_splits, \n        shuffle=True, \n        random_state=CFG.seed).split(np.arange(train.shape[0]), train.label.values)\n    \n    cross_accuracy = []\n    for fold,(train_idx,val_idx) in enumerate(folds):\n        ########\n        # load data\n        #######\n        train_loader,val_loader = load_dataloader(CFG, train, train_idx, val_idx)\n        \n        device = torch.device(CFG.device)\n#         assert(CFG.num_classes ==  train.label.nunique())\n        model = CassavaClassifier(CFG.arch, train.label.nunique(), pretrained=True).to(device)\n        \n        scaler = GradScaler()\n        optimizer = torch.optim.Adam(\n            model.parameters(), \n            lr=CFG.lr, \n            weight_decay=CFG.weight_decay)\n\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n            optimizer, \n            T_0=CFG.T_0, \n            T_mult=CFG.T_mult, \n            eta_min=CFG.min_lr, \n            last_epoch=-1)\n    \n        ########\n        # criterion\n        #######\n        if CFG.criterion == 'LabelSmoothingCrossEntropy':  #### label smoothing cross entropy\n            loss_train = LabelSmoothingCrossEntropy(smoothing=CFG.label_smoothing)\n        else:\n            loss_train = nn.CrossEntropyLoss().to(device)\n        loss_val = nn.CrossEntropyLoss().to(device)\n        \n        best_accuracy = 0\n        best_epoch = 0\n        for epoch in range(CFG.epochs):\n            if epoch < CFG.freeze_bn_epochs:\n                freeze_batchnorm_stats(model)  \n            train_one_epoch(\n                epoch, \n                model, \n                loss_train, \n                optimizer, \n                train_loader, \n                device, \n                scheduler=scheduler, \n                schd_batch_update=False)\n\n            with torch.no_grad():\n                epoch_accuracy = valid_one_epoch(\n                    epoch, \n                    model, \n                    loss_val, \n                    val_loader, \n                    device, \n                    scheduler=None, \n                    schd_loss_update=False)\n\n            if epoch_accuracy > best_accuracy:\n                torch.save(model.state_dict(),'{}_fold{}_best.ckpt'.format(CFG.arch, fold))\n                best_accuracy = epoch_accuracy\n                best_epoch = epoch\n                print('Best model is saved')\n        cross_accuracy += [best_accuracy]\n        print('Fold{} best accuracy = {} in epoch {}'.format(fold,best_accuracy,best_epoch))\n        del model, optimizer, train_loader, val_loader, scaler, scheduler\n        torch.cuda.empty_cache()\n    print('{} folds cross validation CV = {:.5f}'.format(CFG.num_splits,np.average(cross_accuracy)))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Please Upvote if you liked the kernel ! Cheers.\n\nReferences :\nhttps://www.kaggle.com/khyeh0719/pytorch-efficientnet-baseline-train-amp-aug. \nPlease Upvote too."}],"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}