{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\npackage_path = ['../input/timmpackagelatestwhl', '../input/vistion-transformer-pytorch/jx_vit_base_p16_224-80ecf9dd.pth']\nfor pth in package_path:\n    sys.path.append(pth)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport pandas as pd\npd.set_option('display.max_row', None)\npd.set_option('display.max_columns', None)\nimport albumentations as albu\nimport matplotlib.pyplot as plt\nimport json\nimport seaborn as sns\nimport cv2\nimport albumentations as albu\nimport numpy as np\nimport random\n\nimport torch\nimport torch.nn as nn\nimport torchvision.models as models\nfrom torch.optim import Adam\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau, CosineAnnealingWarmRestarts\nfrom sklearn.metrics import accuracy_score, f1_score\nfrom sklearn.model_selection import StratifiedKFold\nfrom albumentations.pytorch import ToTensorV2\nfrom torch.cuda.amp import autocast, GradScaler\n\nfrom tqdm import tqdm\n\n!pip install ../input/timmpackagelatestwhl/timm-0.3.4-py3-none-any.whl\nimport timm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load TrainSet"},{"metadata":{"trusted":true},"cell_type":"code","source":"BASE_DIR=\"../input/cassava-leaf-disease-classification/\"\nTRAIN_IMAGES_DIR=os.path.join(BASE_DIR,'train_images')\ntrain_df=pd.read_csv(os.path.join(BASE_DIR,'train.csv'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"display(train_df.head())\nprint(train_df.shape)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data Loader and Augmentation"},{"metadata":{"trusted":true},"cell_type":"code","source":"class CassavaDataset(Dataset):\n    def __init__(self,df:pd.DataFrame,imfolder:str,train:bool = True, transforms=None):\n        self.df=df\n        self.imfolder=imfolder\n        self.train=train\n        self.transforms=transforms\n        \n    def __getitem__(self,index):\n        im_path=os.path.join(self.imfolder,self.df.iloc[index]['image_id'])\n        x=cv2.imread(im_path,cv2.IMREAD_COLOR)\n        x=cv2.cvtColor(x,cv2.COLOR_BGR2RGB)\n        \n        if(self.transforms):\n            x=self.transforms(image=x)['image']\n        \n        if(self.train):\n            y=self.df.iloc[index]['label']\n            return x,y\n        else:\n            return x\n        \n    def __len__(self):\n        return len(self.df)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_augs = albu.Compose([\n    albu.RandomResizedCrop(height=224, width=224, p=1.0),\n    albu.HorizontalFlip(p=0.5),\n    albu.VerticalFlip(p=0.5),\n    albu.ShiftScaleRotate(p=0.5),\n    albu.Normalize(    \n        mean=[0.485, 0.456, 0.406],\n        std=[0.229, 0.224, 0.225],),\n    ToTensorV2(),\n])\n\nvalid_augs = albu.Compose([\n    albu.Resize(height=224, width=224, p=1.0),\n    albu.Normalize(\n        mean=[0.485, 0.456, 0.406],\n        std=[0.229, 0.224, 0.225],),\n    ToTensorV2(),\n])","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\ndef save_model(model, optimizer, scheduler, fold, epoch, save_every=False, best=False):\n    state = {\n        'model': model.state_dict(),\n        'optimizer': optimizer.state_dict(),\n        'scheduler': scheduler.state_dict()\n    }\n    if save_every == True:\n        if not (os.path.isdir('./saved_model')): os.mkdir('./saved_model')\n        torch.save(state, './saved_model/model_fold_{}_epoch_{}'.format(fold+1, epoch+1))\n    if best == True:\n        if not (os.path.isdir('./best_model')): os.mkdir('./best_model')\n        torch.save(state, './best_model/model_fold_{}_epoch_{}'.format(fold+1, epoch+1))\n        \ndef data_loader(dataset, batch_size, num_workers, phase='train'):\n    if phase == 'train':\n        dataloader = DataLoader(dataset=dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True)\n    else: # valid, test\n        dataloader = DataLoader(dataset=dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False)\n    return dataloader\n        \nclass EarlyStopping:\n    def __init__(self, patience):\n        self.patience = patience\n        self.counter = 0\n        self.early_stop = False\n        self.val_loss_min = np.Inf\n\n    def __call__(self, val_loss, model, optimizer, scheduler, fold, epoch):\n        if self.val_loss_min == np.Inf:\n            self.val_loss_min = val_loss\n        elif val_loss > self.val_loss_min:\n            self.counter += 1\n            print('EarlyStopping counter: {} out of {}'.format(self.counter, self.patience))\n            if self.counter >= self.patience:\n                print('Early Stopping - Fold {} Training is Stopping'.format(fold))\n                self.early_stop = True\n        else:  # val_loss < val_loss_min\n            save_model(model, optimizer, scheduler, fold, epoch, best=True)\n            print('*** Validation loss decreased ({} --> {}).  Saving model... ***'.\\\n                  format(round(self.val_loss_min, 6), round(val_loss, 6)))\n            self.val_loss_min = val_loss\n            self.counter = 0","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Create Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Model(nn.Module):\n    def __init__(self, model_name, n_class, pretrained=False):\n        super().__init__()\n        self.model = timm.create_model(model_name, pretrained=pretrained)\n        n_features = self.model.head.in_features\n        self.model.head = 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":"# train / Validation Functions"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_one_epoch(epoch, model, loss_fn, optimizer, train_loader, device, scheduler):\n    model.train()\n    lst_out = []\n    lst_label = []\n    avg_loss = 0\n\n    status = tqdm(enumerate(train_loader), total=len(train_loader))\n    for step, (images, labels) in status:\n        images = images.to(device).float()\n        labels = labels.to(device).long()\n        with autocast():\n            preds = model(images)\n            lst_out += preds.argmax(1)\n            lst_label += labels\n\n            loss = loss_fn(preds, labels)\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad()\n\n            avg_loss += loss.item() / len(train_loader)\n    scheduler.step()\n    accuracy = accuracy_score(y_pred=torch.tensor(lst_out), y_true=torch.tensor(lst_label))\n    print('{} epoch - train loss : {}, train f1 score : {}'.\\\n          format(epoch + 1, np.round(avg_loss,6), np.round(accuracy*100,2)))\n\ndef valid_one_epoch(epoch, model, loss_fn, val_loader, device):\n    model.eval()\n    lst_val_out = []\n    lst_val_label = []\n    avg_val_loss = 0\n    status = tqdm(enumerate(val_loader), total=len(val_loader))\n    for step, (images, labels) in enumerate(val_loader): #status\n        val_images = images.to(device).float()\n        val_labels = labels.to(device).long()\n\n        val_preds = model(val_images)\n        lst_val_out += val_preds.argmax(1)\n        lst_val_label += val_labels\n        loss = loss_fn(val_preds, val_labels)\n        avg_val_loss += loss.item() / len(val_loader)\n\n    accuracy = accuracy_score(y_pred=torch.tensor(lst_val_out), y_true=torch.tensor(lst_val_label))\n    print('{} epoch - valid loss : {}, valid accuracy : {}'.\\\n          format(epoch + 1, np.round(avg_val_loss, 6), np.round(accuracy*100,2)))\n    return avg_val_loss","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Main - Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"if __name__ == '__main__':\n    train_batch = 16\n    valid_batch = 32\n    num_workers = 4\n    seed = 42\n    split = 5\n    epochs = 100\n    patience = 5\n\n    n_class = 5\n    model_arch = 'vit_base_patch16_224' # 'resnext50_32x4d', 'tf_efficientnet_b4_ns', 'vit_base_patch16_224'\n    device = 'cuda'\n\n    seed_everything(seed)\n    X_train = train_df.iloc[:, :-1]; Y_train = train_df.iloc[:, -1]\n    cv = StratifiedKFold(n_splits=split, random_state=seed, shuffle=True)\n    for fold, (train_index, val_index) in enumerate(cv.split(X_train, Y_train)):\n        print('---------- Fold {} is training ----------'.format(fold + 1))\n        train_x, train_y = X_train.iloc[train_index], Y_train[train_index]\n        val_x, val_y = X_train.iloc[val_index], Y_train[val_index]\n\n        train_dataset=CassavaDataset(df=pd.concat([train_x, train_y], axis=1), imfolder=TRAIN_IMAGES_DIR, train=True, transforms=train_augs)\n        valid_dataset=CassavaDataset(df=pd.concat([val_x, val_y], axis=1), imfolder=TRAIN_IMAGES_DIR, train=True, transforms=valid_augs)\n        train_loader = data_loader(train_dataset, train_batch, num_workers, phase='train')\n        valid_loader = data_loader(valid_dataset, valid_batch, num_workers, phase='valid')\n\n        model = Model(model_arch, n_class, pretrained=False).to(device)\n        model.load_state_dict(torch.load('../input/vistion-transformer-pytorch/jx_vit_base_p16_224-80ecf9dd.pth'), strict=False)\n        loss_tr = nn.CrossEntropyLoss().to(device); loss_fn = nn.CrossEntropyLoss().to(device)\n        optimizer = Adam([param for param in model.parameters() if param.requires_grad], lr=0.001, weight_decay=1e-4)\n        scaler = GradScaler()\n        scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=1, eta_min=1e-6, last_epoch=-1)\n        early_stopping = EarlyStopping(patience=patience)\n\n        for epoch in range(epochs):\n            train_one_epoch(epoch, model, loss_tr, optimizer, train_loader, device, scheduler=scheduler)\n            save_model(model, optimizer, scheduler, fold, epoch)\n            with torch.no_grad():\n                val_loss = valid_one_epoch(epoch, model, loss_fn, valid_loader, device)\n                early_stopping(val_loss, model, optimizer, scheduler, fold, epoch)\n                if early_stopping.early_stop:\n                    break\n\n        del model, optimizer, train_dataset, val_dataset, 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}