{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\n\nOUTPUT_DIR = \"./\"\nTEST_PATH = '../input/cassava-leaf-disease-classification/test_images'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\nsys.path.append('../input/pytorchimagemodels/')\nsys.path.append('../input/pretrainedmodels/')\nsys.path.append('../input/facebook/')\n\nimport time\nimport random\nfrom functools import partial\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\n\nimport cv2\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nfrom torch.optim import Adam\nimport torch.nn.functional as F\nfrom torch.nn import BCEWithLogitsLoss, CrossEntropyLoss\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau, LambdaLR\nfrom torch.utils.data import DataLoader, Dataset\nfrom albumentations import (\n    Compose, OneOf, Normalize, Resize, RandomResizedCrop, RandomCrop, HorizontalFlip, VerticalFlip, \n    RandomBrightness, RandomContrast, RandomBrightnessContrast, Rotate, ShiftScaleRotate, Cutout, \n    IAAAdditiveGaussianNoise, Transpose, HueSaturationValue, \n    )\nfrom albumentations.pytorch import ToTensorV2\n\nimport timm\nfrom timm.models.vision_transformer import VisionTransformer\nfrom pretrainedmodels import se_resnext101_32x4d\nfrom models import DistilledVisionTransformer\n\nimport warnings \nwarnings.filterwarnings('ignore')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_torch(seed=1006):\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    \nseed_torch()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test = pd.read_csv('../input/cassava-leaf-disease-classification/sample_submission.csv')\ntest.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TestDataset(Dataset):\n    def __init__(self, df, transform=None):\n        self.df = df\n        self.file_names = df['image_id'].values\n        self.transform = transform\n        \n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        file_name = self.file_names[idx]\n        file_path = f'{TEST_PATH}/{file_name}'\n        image = cv2.imread(file_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented['image']\n        return image","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_transforms(*, data, vit=False):\n    \n    if vit:\n        MEAN = [0.5, 0.5, 0.5]\n        STD = [0.5, 0.5, 0.5]\n    else:\n        MEAN = [0.485, 0.456, 0.406]\n        STD = [0.229, 0.224, 0.225]\n    \n    if data == 'train':\n        return Compose([\n            RandomResizedCrop(IMG_SIZE, IMG_SIZE),\n            #RandomCrop(IMG_SIZE, 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            #Resize(IMG_SIZE, IMG_SIZE),\n            Normalize(\n                mean=MEAN,\n                std=STD,\n            ),\n            ToTensorV2(),\n        ])\n\n    elif data == 'valid':\n        return Compose([\n            Resize(IMG_SIZE, IMG_SIZE),\n            Normalize(\n                mean=MEAN,\n                std=STD,\n            ),\n            ToTensorV2(),\n        ])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class NetVit(nn.Module):\n    def __init__(self, model_name, pretrained=False, n_class=5, att_activate=False, no_att=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.Identity()\n        \n        if att_activate:\n            self.att_layer = nn.Sequential(\n                nn.Linear(n_features, 256),\n                nn.Tanh(),\n                nn.Linear(256, 1),\n            )\n        else:\n            if no_att:\n                pass\n            else:\n                self.att_layer = nn.Linear(n_features, 1)\n            \n        self.head = nn.Linear(n_features, n_class)\n\n    def forward(self, x):\n        x = self.model(x)\n        output = self.head(x)\n        return output","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class NetVit4(nn.Module):\n    def __init__(self, model_name, pretrained=False, n_class=5, att_activate=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.Identity()\n        if att_activate:\n            self.att_layer = nn.Sequential(\n                nn.Linear(n_features, 256),\n                nn.Tanh(),\n                nn.Linear(256, 1),\n            )\n        else:\n            self.att_layer = nn.Linear(n_features, 1)\n            \n        self.head = nn.Linear(n_features, n_class)\n\n    def forward(self, x):\n        l = x.shape[2] // 2\n        h1 = self.model(x[:, :, :l, :l])\n        h2 = self.model(x[:, :, :l, l:])\n        h3 = self.model(x[:, :, l:, :l])\n        h4 = self.model(x[:, :, l:, l:])\n\n        a1 = self.att_layer(h1)\n        a2 = self.att_layer(h2)\n        a3 = self.att_layer(h3)\n        a4 = self.att_layer(h4)\n\n        w = F.softmax(torch.cat([a1, a2, a3, a4], dim=1), dim=1)\n\n        h = h1 * w[:, 0].unsqueeze(-1) + h2 * w[:, 1].unsqueeze(-1) + h3 * w[:, 2].unsqueeze(-1) + h4 * w[:, 3].unsqueeze(-1)\n        output = self.head(h)\n        return output","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from collections import OrderedDict\n\ndef inference(model, states, test_loader, device, temp=1):\n    model.to(device)\n    preds = []\n    for state in states:\n        pred = []\n        model.load_state_dict(state)\n        model.eval()\n        for i, image in enumerate(test_loader):\n            with torch.no_grad():\n                pred.append((model(image.to(device))*temp).softmax(1).to('cpu'))\n        pred = torch.cat(pred, dim=0)\n        preds.append(pred.numpy())\n    return np.mean(preds, axis=0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def multi2single(path, se=False):\n    state_dict = torch.load(path, map_location=lambda storage, loc: storage)\n    new_state_dict = OrderedDict()\n    for k, v in state_dict.items():\n        if 'module' in k:\n            k = k.replace('se_module', 'dummy')\n            k = k.replace('module.', '')\n            k = k.replace('dummy', 'se_module')\n        if 'attention_linear' in k:\n            k = k.replace('attention_linear', 'att_layer')\n        new_state_dict[k] = v\n    return new_state_dict","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"temp = 1.0","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"    MODEL_NAME = \"vit_base_patch16_384\"\n    MODEL_NUM = \"No3001\"\n    MODEL_DIR = \"../input/cassavamodels/\"\n    IMG_SIZE = 384\n    TTA = 5\n    BATCH = 32\n\n    model = NetVit(MODEL_NAME, pretrained=False, no_att=True)\n    states = [multi2single(MODEL_DIR+f'{MODEL_NUM}_{fold+1}.pth') for fold in range(5)]\n    if TTA == 1:\n        test_dataset = TestDataset(test, transform=get_transforms(data='valid', vit=True))\n    else:\n        test_dataset = TestDataset(test, transform=get_transforms(data='train', vit=True))\n\n    test_loader = DataLoader(test_dataset, batch_size=BATCH, shuffle=False, num_workers=4, pin_memory=True)\n    vit_predictions = np.zeros((len(test), 5))\n    for _ in range(TTA):\n        vit_predictions += inference(model, states, test_loader, device, temp) / TTA","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if True:\n    MODEL_NAME = \"vit_base_patch16_224\"\n    MODEL_NUM = \"vit4_ex\"\n    MODEL_DIR = \"../input/cassavamodels/\"\n    IMG_SIZE = 448\n    TTA = 5\n    BATCH = 32\n    att_activate = False\n\n    model = NetVit4(MODEL_NAME, pretrained=False, att_activate=att_activate)    \n    states = [multi2single(MODEL_DIR+f'{MODEL_NUM}_{fold+1}.pth') for fold in range(5)]\n    if TTA == 1:\n        test_dataset = TestDataset(test, transform=get_transforms(data='valid', vit=True))\n    else:\n        test_dataset = TestDataset(test, transform=get_transforms(data='train', vit=True))\n\n    test_loader = DataLoader(test_dataset, batch_size=BATCH, shuffle=False, num_workers=4, pin_memory=True)\n    vit4_predictions_a = np.zeros((len(test), 5))\n    for _ in range(TTA):\n        vit4_predictions_a += inference(model, states, test_loader, device, temp) / TTA","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if True:\n    MODEL_NAME = \"vit_base_patch16_224\"\n    MODEL_NUM = \"vit4_ex_smooth001_att_act\"\n    MODEL_DIR = \"../input/cassavamymodels/\"\n    IMG_SIZE = 448\n    TTA = 5\n    BATCH = 32\n    att_activate = True\n\n    model = NetVit4(MODEL_NAME, pretrained=False, att_activate=att_activate)    \n    states = [multi2single(MODEL_DIR+f'{MODEL_NUM}_{fold+1}.pth') for fold in range(5)]\n    if TTA == 1:\n        test_dataset = TestDataset(test, transform=get_transforms(data='valid', vit=True))\n    else:\n        test_dataset = TestDataset(test, transform=get_transforms(data='train', vit=True))\n\n    test_loader = DataLoader(test_dataset, batch_size=BATCH, shuffle=False, num_workers=4, pin_memory=True)\n    vit4_predictions_b = np.zeros((len(test), 5))\n    for _ in range(TTA):\n        vit4_predictions_b += inference(model, states, test_loader, device, temp) / TTA","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"predictions = (vit_predictions * 0.45 + vit4_predictions_a * 0.55) / 9 * 10 + vit4_predictions_b * 0.08","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# submission\ntest['label'] = predictions.argmax(1)\ntest[['image_id', 'label']].to_csv(OUTPUT_DIR+'submission.csv', index=False)\ntest.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}