{"cells":[{"metadata":{},"cell_type":"markdown","source":"- Reference <br>\nhttps://www.kaggle.com/szuzhangzhi/vision-transformer-vit-cuda-as-usual/data\n\nI made a notebook by dividing it into [Train] and [Inference] parts by referring to the original notebook.\n\n### Please upvote original notebook :)"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# !pip install vision_transformer_pytorch\n# !pip install AdamP","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\n\npackage_path = '../input/vision-transformer-pytorch/VisionTransformer-Pytorch'\nsys.path.append(package_path)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Library"},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport pandas as pd\n\nimport time\nimport datetime\nimport copy\nimport matplotlib.pyplot as plt\nimport json\nimport seaborn as sns\nimport cv2\nimport albumentations as albu\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nimport torchvision.models as models\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom sklearn.metrics import accuracy_score\nfrom sklearn.model_selection import StratifiedKFold, GroupKFold, KFold, train_test_split\n\n\n# ALBUMENTATIONS\nimport albumentations as albu\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)\n    \nfrom albumentations.pytorch import ToTensorV2\n\n# ADAMP\n# from adamp import AdamP","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"BASE_DIR=\"../input/cassava-leaf-disease-classification/\"\nTRAIN_IMAGES_DIR = os.path.join(BASE_DIR,'train_images')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(os.path.join(BASE_DIR, 'train.csv'))\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(\"Count of training images {0}\".format(len(os.listdir(TRAIN_IMAGES_DIR))))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with open(f'{BASE_DIR}/label_num_to_disease_map.json', 'r') as f:\n    name_mapping = json.load(f)\n    \nname_mapping = {int(k): v for k, v in name_mapping.items()}\ntrain_df['class_id'] = train_df['label'].map(name_mapping)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"name_mapping","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visualization"},{"metadata":{"trusted":true},"cell_type":"code","source":"def visualize_images(image_ids, labels):\n    plt.figure(figsize=(16, 12))\n    \n    for idx, (image_id, label) in enumerate(zip(image_ids, labels)):\n        plt.subplot(3, 3, idx+1)\n        \n        image = cv2.imread(os.path.join(TRAIN_IMAGES_DIR, image_id))\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        plt.imshow(image)\n        plt.title(f\"Class: {label}\", fontsize=12)\n        \n        plt.axis(\"off\")\n        \n    plt.show()\n    \n\ndef plot_augmentation(image_id, transform):\n    plt.figure(figsize=(16, 4))\n    \n    img = cv2.imread(os.path.join(TRAIN_IAMGES_DIR, image_id))\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    \n    plt.subplot(1, 3, 1)\n    plt.imshow(img)\n    plt.axis('off')\n    \n    plt.subplot(1, 3, 2)\n    x = transform(image=img)['image']\n    plt.imshow(x)\n    plt.axis('off')\n    \n    plt.subplot(1, 3, 3)\n    x = transform(image=img)['image']\n    plt.imshow(x)\n    \ndef visualize(images, transform):\n    '''\n    Plot images and their transformations\n    '''\n    fig = plt.figure(figsize=(32, 16))\n    \n    for i, im in enumerate(images):\n        ax = fig.add_subplot(2, 5, i+1, xticks=[], yticks=[])\n        plt.imshow(im)\n        \n    for i, im in enumerate(images):\n        ax = fig.add_subplot(2, 5, i+6, xticks=[], yticks=[])\n        plt.imshow(transform(image=im)['image'])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Modeling"},{"metadata":{"trusted":true},"cell_type":"code","source":"# CUSTOM DATASET CLASS\nclass CassavaDataset(Dataset):\n    def __init__(\n        self, df:pd.DataFrame, imfolder:str, train:bool=True, transforms=None\n    ):\n        self.df = df\n        self.imfolder = imfolder\n        self.train = train\n        self.transforms = transforms\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, index):\n        im_path = os.path.join(self.imfolder, self.df.iloc[index]['image_id'])\n        im = cv2.imread(im_path, cv2.IMREAD_COLOR)\n        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)\n        \n        if (self.transforms):\n            '''\n            When AlbumentationCompose, a dictionary with key 'image' is created\n            '''\n            im = self.transforms(image=im)['image']\n            \n        if (self.train):\n            label = self.df.iloc[index]['label']\n            return im, label\n        else:\n            return im","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# AUGMENTATIONS\ntrain_augs = albu.Compose([\n    albu.RandomResizedCrop(height=384, width=384, p=1.0),\n    albu.HorizontalFlip(p=0.5),\n    albu.VerticalFlip(p=0.5),\n    albu.RandomBrightnessContrast(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    CoarseDropout(p=0.5),\n    Cutout(p=0.5),\n    ToTensorV2(),\n])\n\nvalid_augs = albu.Compose([\n    albu.Resize(height=384, width=384, 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":{"trusted":true},"cell_type":"code","source":"# DATA SPLIT\ntrain, valid = train_test_split(\n    train_df,\n    test_size=0.1,\n    random_state=42,\n    stratify=train_df.label.values\n)\n\n# reset index on both dataframes\ntrain = train.reset_index(drop=True)\nvalid = valid.reset_index(drop=True)\n\n# targets in train,valid datasets\ntrain_targets = train.label.values\nvalid_targets = valid.label.values","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# DEFINE PYTORCH CUSTOM DATASET\ntrain_dataset = CassavaDataset(\n    df = train,\n    imfolder = TRAIN_IMAGES_DIR,\n    train = True,\n    transforms = train_augs\n)\n\nvalid_dataset = CassavaDataset(\n    df = valid,\n    imfolder = TRAIN_IMAGES_DIR,\n    train=True,\n    transforms = valid_augs\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_image(img_dict):\n    image_tensor = img_dict[0]\n#     print(type(image_tensor))\n    target = img_dict[1]\n    print(target)\n    plt.figure(figsize=(10, 10))\n    image = image_tensor.permute(1, 2, 0)\n    plt.imshow(image)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plot_image(train_dataset[5])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# MAKE PYTORCH DATALOADER\ntrain_loader = DataLoader(\n    train_dataset,\n    batch_size = 16,\n    num_workers = 4,\n    shuffle = True\n)\n\nvalid_loader = DataLoader(\n    valid_dataset,\n    batch_size = 16,\n    num_workers = 4,\n    shuffle = False\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# TRAIN\ndef train_model(datasets, dataloaders, model, criterion, optimizer, scheduler, num_epochs, device):\n    since = time.time()\n    \n    best_model_wts = copy.deepcopy(model.state_dict())\n    best_acc = 0.0\n    \n    for epoch in range(num_epochs):\n        print('Epoch {}/{}'.format(epoch, num_epochs-1))\n        print('-' * 10)\n        \n        for phase in ['train', 'valid']:\n            if phase == 'train':\n                model.train()\n            else:\n                model.eval()\n                \n            running_loss = 0.0\n            running_corrects = 0.0\n            \n            for inputs, labels in dataloaders[phase]:\n                inputs = inputs.to(device)\n                labels = labels.to(device)\n                \n                # Zero out the grads\n                optimizer.zero_grad()\n                \n                # Forward\n                # Track history in train mode\n                with torch.set_grad_enabled(phase == 'train'):\n                    model = model.to(device)\n                    outputs = model(inputs)\n                    _, preds = torch.max(outputs, 1) \n                    loss = criterion(outputs, labels)\n                    \n                    if phase == 'train':\n                        loss.backward()\n                        optimizer.step()\n                \n                # Statistics\n                running_loss += loss.item()*inputs.size(0)\n                running_corrects += torch.sum(preds == labels.data)\n            \n            if phase == 'train':\n                scheduler.step()\n                \n            epoch_loss = running_loss / len(datasets[phase])\n            epoch_acc = running_corrects.double() / len(datasets[phase])\n            \n            print('{} Loss: {:.4f} Acc: {:.4f}'.format(\n                phase, epoch_loss, epoch_acc))\n            \n            if phase == 'valid' and epoch_acc > best_acc:\n                best_acc = epoch_acc\n                best_model_wts = copy.deepcopy(model.state_dict())\n        \n        print()\n    \n    time_elapsed = time.time() - since\n    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n    print('Best val Acc: {:.4f}'.format(best_acc))\n    \n    model.load_state_dict(best_model_wts)\n    \n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from vision_transformer_pytorch import VisionTransformer\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ndatasets = {'train': train_dataset,\n            'valid': valid_dataset}\n\ndataloaders = {'train': train_loader,\n               'valid': valid_loader}\n\n# LOAD PRETRAINED ViT MODEL\nmodel = VisionTransformer.from_pretrained('ViT-B_16', num_classes=5)        \n\n# OPTIMIZER\noptimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.001)\n# optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=0.001)\n# optimizer = AdamP(model.parameters(), lr=1e-4, weight_decay=0.001)\n\n# LEARNING RATE SCHEDULER\nscheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.1)\n\ncriterion = nn.CrossEntropyLoss()\nnum_epochs = 6","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# MODEL TRAIN\ntrained_model = train_model(datasets, dataloaders, model, criterion, optimizer, scheduler, num_epochs, device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Save the mode after training\ntorch.save(model.state_dict(), 'vit_b-16.pt')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Inference: https://www.kaggle.com/nevret93/cassava-vit-cuda-inference"}],"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}