{"cells":[{"metadata":{},"cell_type":"markdown","source":"# [PyTorch] Attention Module on Local Features\n\n\nIn addition to the usual global average pooling, this notebook illustrate __an idea__ of adding an existing visual attention (channel and spatial) module for local features modelling. The attention module is the attention module from [DANet](https://arxiv.org/abs/1809.02983), implementation is their official implementation, adapted for this special used case.\n\nWith the attention module, the __boost__ to local CV is around 0.05 (~0.945 -> 0.950), LB is around 0.04 (0.942 (V2) -> 0.946 (V1)). All in a single-fold (out of 5 folds).\n\nIn this notebook, the backbone is efficientnet-b2, but the method can be adapted to larger backbone like the popular resnet200d or efficientnet-b7.\n\nP.S: In the process of trying this method, I realized the quality of training has a huge impact on the model performance in this competition (most Kaggler probably know this already lol....)."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport sys\nimport os\nimport glob\n\nimport cv2\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport albumentations as a_transform\nfrom torch.utils.data import DataLoader, Dataset\n\nfrom albumentations.pytorch import ToTensorV2\neffnet_pytorch = \"../input/efficientnet-pytorch/EfficientNet-PyTorch/EfficientNet-PyTorch-master\"\nsys.path.append(effnet_pytorch)\nfrom efficientnet_pytorch import EfficientNet\n\ndevice = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\nprint(device)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Attention Module\n\nSimilar idea with CBAM (inspiration from this [kernel](https://www.kaggle.com/ipythonx/tf-keras-ranzcr-multi-attention-efficientnet)) and [Dual Attention network](https://github.com/junfu1115/DANet)."},{"metadata":{"trusted":true},"cell_type":"code","source":"class PAM_Module(nn.Module):\n    \"\"\" Position attention module\"\"\"\n    #Ref from SAGAN\n    def __init__(self, in_dim):\n        super(PAM_Module, self).__init__()\n        self.chanel_in = in_dim\n\n        self.query_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim//8, kernel_size=1)\n        self.key_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim//8, kernel_size=1)\n        self.value_conv = nn.Conv2d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)\n        self.gamma = nn.Parameter(torch.zeros(1))\n\n    def forward(self, x):\n        \"\"\"\n            inputs :\n                x : input feature maps( B X C X H X W)\n            returns :\n                out : attention value + input feature\n                attention: B X (HxW) X (HxW)\n        \"\"\"\n        m_batchsize, C, height, width = x.size()\n        proj_query = self.query_conv(x).view(m_batchsize, -1, width*height).permute(0, 2, 1)\n        proj_key = self.key_conv(x).view(m_batchsize, -1, width*height)\n        energy = torch.bmm(proj_query, proj_key)\n        attention = torch.softmax(energy, dim=-1)\n        proj_value = self.value_conv(x).view(m_batchsize, -1, width*height)\n\n        out = torch.bmm(proj_value, attention.permute(0, 2, 1))\n        out = out.view(m_batchsize, C, height, width)\n\n        out = self.gamma*out + x\n        return out\n\n\nclass CAM_Module(nn.Module):\n    \"\"\" Channel attention module\"\"\"\n    def __init__(self, in_dim):\n        super(CAM_Module, self).__init__()\n        self.chanel_in = in_dim\n        self.gamma = nn.Parameter(torch.zeros(1))\n\n    def forward(self,x):\n        \"\"\"\n            inputs :\n                x : input feature maps( B X C X H X W)\n            returns :\n                out : attention value + input feature\n                attention: B X C X C\n        \"\"\"\n        m_batchsize, C, height, width = x.size()\n        proj_query = x.view(m_batchsize, C, -1)\n        proj_key = x.view(m_batchsize, C, -1).permute(0, 2, 1)\n        energy = torch.bmm(proj_query, proj_key)\n        energy_new = torch.max(energy, -1, keepdim=True)[0].expand_as(energy)-energy\n        attention = torch.softmax(energy_new, dim=-1)\n        proj_value = x.view(m_batchsize, C, -1)\n\n        out = torch.bmm(attention, proj_value)\n        out = out.view(m_batchsize, C, height, width)\n\n        out = self.gamma*out + x\n        return out\n\n\nclass CBAM(nn.Module):\n    def __init__(self, in_channels):\n        # def __init__(self):\n        super(CBAM, self).__init__()\n        inter_channels = in_channels // 4\n        self.conv1_c = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),\n                                     nn.BatchNorm2d(inter_channels),\n                                     nn.ReLU())\n        \n        self.conv1_s = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),\n                                     nn.BatchNorm2d(inter_channels),\n                                     nn.ReLU())\n\n        self.channel_gate = CAM_Module(inter_channels)\n        self.spatial_gate = PAM_Module(inter_channels)\n\n        self.conv2_c = nn.Sequential(nn.Conv2d(inter_channels, in_channels, 3, padding=1, bias=False),\n                                     nn.BatchNorm2d(in_channels),\n                                     nn.ReLU())\n        self.conv2_a = nn.Sequential(nn.Conv2d(inter_channels, in_channels, 3, padding=1, bias=False),\n                                     nn.BatchNorm2d(in_channels),\n                                     nn.ReLU())\n\n    def forward(self, x):\n        feat1 = self.conv1_c(x)\n        chnl_att = self.channel_gate(feat1)\n        chnl_att = self.conv2_c(chnl_att)\n\n        feat2 = self.conv1_s(x)\n        spat_att = self.spatial_gate(feat2)\n        spat_att = self.conv2_a(spat_att)\n\n        x_out = chnl_att + spat_att\n\n        return x_out\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# EfficientNet-b2 with soft attention"},{"metadata":{"trusted":true},"cell_type":"code","source":"class EffNetWLF(nn.Module):\n\n    def __init__(self, model_name, target_size=11):\n        super().__init__()\n        self.backbone = EfficientNet.from_name(model_name)\n\n        self.backbone._dropout = nn.Dropout(0.1)\n        n_features = self.backbone._fc.in_features\n        self.backbone._fc = nn.Linear(n_features, target_size)\n\n        self.local_fe = CBAM(n_features)\n        self.dropout = nn.Dropout(0.1)\n        self.classifier = nn.Sequential(nn.Linear(n_features + n_features, n_features),\n                                        nn.BatchNorm1d(n_features),\n                                        nn.Dropout(0.1),\n                                        nn.ReLU(),\n                                        nn.Linear(n_features, target_size))\n\n    def forward(self, image):\n        enc_feas = self.backbone.extract_features(image)\n\n        # use default's global features\n        global_feas = self.backbone._avg_pooling(enc_feas)\n        global_feas = global_feas.flatten(start_dim=1)\n        global_feas = self.dropout(global_feas)\n\n        local_feas = self.local_fe(enc_feas)\n        local_feas = torch.sum(local_feas, dim=[2,3])\n        local_feas = self.dropout(local_feas)\n\n        all_feas = torch.cat([global_feas, local_feas], dim=1)\n        outputs = self.classifier(all_feas)\n        return outputs\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset and DataLoader"},{"metadata":{"trusted":true},"cell_type":"code","source":"work_dir = \"../input/ranzcr-clip-catheter-line-classification/\"\ndf = pd.read_csv(os.path.join(work_dir, \"sample_submission.csv\"))\ntest_img_paths = glob.glob(os.path.join(work_dir, \"test/*.jpg\"))\nprint(len(test_img_paths))\nmodel_weights = glob.glob(\"../input/effb2wlf/*.pth\")\nprint(model_weights)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class TestDataset(Dataset):\n    def __init__(self, img_paths, transform=None):\n        self.img_paths = img_paths\n        self.transform = transform\n        self.clahe = cv2.createCLAHE(clipLimit=30.0, tileGridSize=(8, 8))\n        \n    def __len__(self):\n        return len(self.img_paths)\n\n    def __getitem__(self, idx):\n        image = cv2.imread(self.img_paths[idx], cv2.IMREAD_GRAYSCALE)\n        image = self.clahe.apply(image)\n        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)\n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented['image']\n        uid = self.img_paths[idx][:-4].split('/')[-1]\n        return uid, image","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataloader"},{"metadata":{"trusted":true},"cell_type":"code","source":"normalize = a_transform.Normalize(mean=[0.485, 0.456, 0.406],\n                                  std=[0.229, 0.224, 0.225], p=1.0, max_pixel_value=255.0)\ntest_transform = a_transform.Compose([a_transform.Resize(512, 512),\n                                      normalize,\n                                      ToTensorV2()], p=1.0)\ntest_ds = TestDataset(test_img_paths, test_transform)\ndataloader = DataLoader(test_ds, batch_size=16, num_workers=4, shuffle=False, drop_last=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"final_pred = np.empty((len(model_weights),len(test_ds), 11), dtype=np.float32)\nfor model_idx, each_w in enumerate(model_weights):\n    print(f\"running idx {model_idx}\")\n    model = EffNetWLF(\"efficientnet-b2\")\n    model = model.to(device)\n    checkpoint = torch.load(f\"{each_w}\")\n    model.load_state_dict(checkpoint[\"model\"])\n    uids = []\n    all_pred = []\n    for idx, (name, x_mb) in enumerate(dataloader):\n        x = x_mb.to(device)\n        x = torch.stack([x,x.flip(-1)],0) # hflip\n        x = x.view(-1, 3, 512, 512)\n        with torch.no_grad():\n            pred = model(x)\n        pred = pred.view(2, x_mb.size(0), -1).mean(0)\n        uids.append(name)\n        all_pred.append(torch.sigmoid(pred))\n    all_pred = torch.cat(all_pred, dim=0).cpu().numpy()\n    final_pred[model_idx] = all_pred\n    \nfinal_pred = final_pred.mean(axis=0)\nprint(final_pred.shape)\n\ntarget_cols = df.columns[1:]\ndf[target_cols] = final_pred\ndf[\"StudyInstanceUID\"] = np.concatenate(uids)\ndf.to_csv('submission.csv', index=False)\ndf.head()","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}