vgg_cbam.ipynb

{"cells":[{"cell_type":"code","execution_count":1,"metadata":{"executionInfo":{"elapsed":2555,"status":"ok","timestamp":1705839708370,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"0fwj2xBjE4E6"},"outputs":[],"source":["import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","import torch.optim as optim\n","import torch.autograd\n","import torch.utils.data\n","import torchvision\n","import torchvision.datasets as dataset\n","import torchvision.transforms as T\n","import matplotlib.pyplot as plt\n","import numpy as np\n","\n","from tqdm import tqdm\n","from torch.utils.data import Dataset\n","from torch.utils.data import DataLoader"]},{"cell_type":"code","execution_count":2,"metadata":{"executionInfo":{"elapsed":6,"status":"ok","timestamp":1705839708370,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"-uNaCfZPgTUu"},"outputs":[],"source":["RANDOM_SEED = 42\n","torch.manual_seed(RANDOM_SEED)\n","torch.cuda.manual_seed(RANDOM_SEED)\n","torch.cuda.manual_seed_all(RANDOM_SEED)\n","torch.backends.cudnn.deterministic = True\n","torch.backends.cudnn.benchmark = False"]},{"cell_type":"markdown","metadata":{"id":"FvfWrKZwE4E8"},"source":["## Dataset"]},{"cell_type":"code","execution_count":3,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":1081,"status":"ok","timestamp":1705839709447,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"AegtXewuc4nw","outputId":"c1ec0d99-7418-40da-9a97-38e98dcd4eca"},"outputs":[{"name":"stdout","output_type":"stream","text":["Files already downloaded and verified\n"]}],"source":["train_dataset = dataset.CIFAR100(root=\"./CIFAR100/train\", train=True, transform=None, download=True)"]},{"cell_type":"code","execution_count":4,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":8371,"status":"ok","timestamp":1705839717807,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"U_guRntscPM9","outputId":"af74b35d-17ad-4492-b373-4a9e244fdfae"},"outputs":[{"name":"stdout","output_type":"stream","text":["mean: [0.5070751592371323, 0.48654887331495095, 0.4409178433670343]\n","std: [0.26733428587941854, 0.25643846292120615, 0.2761504713263903]\n"]}],"source":["x = np.concatenate([np.asarray(train_dataset[i][0]) for i in range(len(train_dataset))])\n","mean = np.mean(x, axis=(0, 1))/255\n","std = np.std(x, axis=(0,1))/255\n","\n","mean = mean.tolist()\n","std = std.tolist()\n","\n","print(\"mean:\", mean)\n","print(\"std:\", std)"]},{"cell_type":"code","execution_count":5,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":2695,"status":"ok","timestamp":1705839720495,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"jnivY6thE4E9","outputId":"4de71cf3-fdcb-41d1-a81c-dfcd87a06752"},"outputs":[{"name":"stdout","output_type":"stream","text":["Files already downloaded and verified\n","Files already downloaded and verified\n","Files already downloaded and verified\n"]}],"source":["transform = T.Compose([T.ToTensor(),\n","                       T.Normalize(mean, std, inplace=True)])\n","train_dataset = dataset.CIFAR100(root=\"./CIFAR100/train\", train=True, transform=transform, download=True)\n","valid_dataset = dataset.CIFAR100(root=\"./CIFAR100/val\", train=True, transform=transform, download=True)\n","test_dataset = dataset.CIFAR100(root=\"./CIFAR100/test\", train=False, transform=transform, download=True)"]},{"cell_type":"code","execution_count":6,"metadata":{"executionInfo":{"elapsed":9,"status":"ok","timestamp":1705839720495,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"-vArjQBOE4E-"},"outputs":[],"source":["train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True, num_workers=8)\n","valid_loader = DataLoader(valid_dataset, batch_size=256, shuffle=False, num_workers=8)\n","test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False, num_workers=8)"]},{"cell_type":"code","execution_count":7,"metadata":{"executionInfo":{"elapsed":384,"status":"ok","timestamp":1705839720871,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"ZIPfRoU6E4E-"},"outputs":[],"source":["images, labels = next(iter(train_loader))"]},{"cell_type":"code","execution_count":8,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":509,"status":"ok","timestamp":1705839721379,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"3nxHFxqmE4E-","outputId":"30e47cfd-a067-4ffe-cd73-ff7d44147774"},"outputs":[{"name":"stdout","output_type":"stream","text":["torch.Size([256, 3, 32, 32])\n","torch.Size([256])\n","tensor([27, 10, 73, 51, 53, 82, 21, 44, 47,  5, 94, 57, 17, 94, 92, 91, 53, 76,\n","        45, 46, 52, 30, 30, 85, 39, 70,  2,  6, 30, 58, 36, 67, 61, 38, 28, 10,\n","        61, 46, 43, 61, 27, 22, 86, 11, 43, 87, 78, 47, 97,  2, 94, 16, 79, 92,\n","        84,  8, 91, 35, 61, 55, 75, 42, 60, 49, 41, 96, 46, 97, 37, 34, 58,  4,\n","        89, 51, 92, 62, 66, 40, 95, 26, 31, 86, 33, 82, 58, 14, 68, 26, 96, 23,\n","        87, 82, 15, 65, 26,  3,  1, 92, 60, 32, 16, 44, 20, 86, 36, 56, 99, 61,\n","        72, 46, 91, 84, 89, 30, 37,  0, 69, 68, 58, 11, 46, 41, 98, 46, 15,  2,\n","        21, 37, 59, 79, 88, 71, 65, 70, 60, 21, 62, 92, 40, 29,  6, 48, 40, 66,\n","        90, 98, 44, 64, 46, 13,  6,  6, 18, 80, 33, 76, 27, 45, 11, 61, 59, 96,\n","        60, 25, 80, 66, 48, 87, 24,  1, 35, 56,  2, 30,  3,  7, 29, 47, 81, 23,\n","         2, 91, 12,  3, 40, 15, 70, 68, 25,  3, 31, 73, 10, 87, 85, 45, 74, 58,\n","         7, 56, 20, 84, 93, 34, 15, 45,  4, 82, 61, 42, 10, 82, 91, 46, 58, 99,\n","        74,  6, 43, 99, 21, 23, 37, 89, 35, 17, 32, 16, 88, 21, 14, 38, 28, 27,\n","        79, 91, 90, 99, 23, 15, 35, 98, 65, 12, 61, 21, 56, 20, 80, 28, 35, 41,\n","        31, 87,  7, 68])\n"]}],"source":["print(images.shape)\n","print(labels.shape)\n","print(labels)"]},{"cell_type":"markdown","metadata":{"id":"2WYBpunpE4E_"},"source":["## Model"]},{"cell_type":"code","execution_count":9,"metadata":{"executionInfo":{"elapsed":10,"status":"ok","timestamp":1705839721380,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"7pCCEZNzE4E_"},"outputs":[],"source":["# No CBAM\n","def conv_block(in_channels, out_channels):\n","    return nn.Sequential(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1),\n","                         nn.BatchNorm2d(out_channels),\n","                         nn.ReLU())\n","\n","class VGG19(nn.Module):\n","    def __init__(self, in_channels, out_channels):\n","        super().__init__()\n","        self.in_channels = in_channels\n","        self.out_channels = out_channels\n","\n","        self.conv_block1 = nn.Sequential(nn.Conv2d(in_channels=self.in_channels, out_channels=64, kernel_size=3, stride=1, padding=1),\n","                                         nn.BatchNorm2d(64),\n","                                         nn.ReLU(),\n","                                         conv_block(64, 64),\n","                                         nn.MaxPool2d(kernel_size=2, stride=2))\n","\n","        self.conv_block2 = nn.Sequential(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),\n","                                         nn.BatchNorm2d(128),\n","                                         nn.ReLU(),\n","                                         conv_block(128, 128),\n","                                         nn.MaxPool2d(kernel_size=2, stride=2))\n","\n","        self.conv_block3 = nn.Sequential(nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),\n","                                         nn.BatchNorm2d(256),\n","                                         nn.ReLU(),\n","                                         *[conv_block(256, 256) for _ in range(3)],\n","                                         nn.MaxPool2d(kernel_size=2, stride=2))\n","\n","        self.conv_block4 = nn.Sequential(nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),\n","                                         nn.BatchNorm2d(512),\n","                                         nn.ReLU(),\n","                                         *[conv_block(512, 512) for _ in range(3)],\n","                                         nn.MaxPool2d(kernel_size=2, stride=2))\n","\n","        self.conv_block5 = nn.Sequential(*[conv_block(512, 512) for _ in range(4)],\n","                                         nn.MaxPool2d(kernel_size=2, stride=2))\n","\n","        self.avg_pool = nn.AdaptiveAvgPool2d(output_size=(7,7))\n","\n","        self.linear1 = nn.Sequential(nn.Linear(in_features=7*7*512, out_features=4096, bias=True),\n","                                     nn.Dropout(0.5),\n","                                     nn.ReLU())\n","        self.linear2 = nn.Sequential(nn.Linear(in_features=4096, out_features=4096, bias=True),\n","                                     nn.Dropout(0.5),\n","                                     nn.ReLU())\n","        self.linear3 = nn.Linear(in_features=4096, out_features=self.out_channels, bias=True)\n","\n","    def forward(self, x):\n","        x = self.conv_block1(x)\n","        x = self.conv_block2(x)\n","        x = self.conv_block3(x)\n","        x = self.conv_block4(x)\n","        x = self.conv_block5(x)\n","        x = self.avg_pool(x)\n","        x = self.linear1(x.view(x.shape[0], -1))\n","        x = self.linear2(x)\n","        x = self.linear3(x)\n","        return x"]},{"cell_type":"code","execution_count":10,"metadata":{"executionInfo":{"elapsed":7,"status":"ok","timestamp":1705839721380,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"Tsc4eye3-Qja"},"outputs":[],"source":["# With CBAM\n","class SAM(nn.Module):\n","    def __init__(self, bias=False):\n","        super(SAM, self).__init__()\n","        self.bias = bias\n","        self.conv = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3, dilation=1, bias=self.bias)\n","\n","    def forward(self, x):\n","        max = torch.max(x,1)[0].unsqueeze(1)\n","        avg = torch.mean(x,1).unsqueeze(1)\n","        concat = torch.cat((max,avg), dim=1)\n","        output = self.conv(concat)\n","        output = output * x\n","        return output\n","\n","class CAM(nn.Module):\n","    def __init__(self, channels, r):\n","        super(CAM, self).__init__()\n","        self.channels = channels\n","        self.r = r\n","        self.linear = nn.Sequential(\n","            nn.Linear(in_features=self.channels, out_features=self.channels//self.r, bias=True),\n","            nn.ReLU(inplace=True),\n","            nn.Linear(in_features=self.channels//self.r, out_features=self.channels, bias=True))\n","\n","    def forward(self, x):\n","        max = F.adaptive_max_pool2d(x, output_size=1)\n","        avg = F.adaptive_avg_pool2d(x, output_size=1)\n","        b, c, _, _ = x.size()\n","        linear_max = self.linear(max.view(b,c)).view(b, c, 1, 1)\n","        linear_avg = self.linear(avg.view(b,c)).view(b, c, 1, 1)\n","        output = linear_max + linear_avg\n","        output = F.sigmoid(output) * x\n","        return output\n","\n","class CBAM(nn.Module):\n","    def __init__(self, channels, r):\n","        super(CBAM, self).__init__()\n","        self.channels = channels\n","        self.r = r\n","        self.sam = SAM(bias=False)\n","        self.cam = CAM(channels=self.channels, r=self.r)\n","\n","    def forward(self, x):\n","        output = self.cam(x)\n","        output = self.sam(output)\n","        return output + x\n","\n","def conv_block(in_channels, out_channels):\n","    return nn.Sequential(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1),\n","                         nn.BatchNorm2d(out_channels),\n","                         nn.ReLU())\n","\n","class VGG19_CBAM(nn.Module):\n","    def __init__(self, in_channels, out_channels):\n","        super().__init__()\n","        self.in_channels = in_channels\n","        self.out_channels = out_channels\n","\n","        self.conv_block1 = nn.Sequential(nn.Conv2d(in_channels=self.in_channels, out_channels=64, kernel_size=3, stride=1, padding=1),\n","                                         nn.BatchNorm2d(64),\n","                                         nn.ReLU(),\n","                                         conv_block(64, 64),\n","                                         CBAM(64, r=4),\n","                                         nn.MaxPool2d(kernel_size=2, stride=2))\n","\n","        self.conv_block2 = nn.Sequential(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),\n","                                         nn.BatchNorm2d(128),\n","                                         nn.ReLU(),\n","                                         conv_block(128, 128),\n","                                         CBAM(128, r=4),\n","                                         nn.MaxPool2d(kernel_size=2, stride=2))\n","\n","        self.conv_block3 = nn.Sequential(nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),\n","                                         nn.BatchNorm2d(256),\n","                                         nn.ReLU(),\n","                                         *[conv_block(256, 256) for _ in range(3)],\n","                                         CBAM(256, r=4),\n","                                         nn.MaxPool2d(kernel_size=2, stride=2))\n","\n","        self.conv_block4 = nn.Sequential(nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1),\n","                                         nn.BatchNorm2d(512),\n","                                         nn.ReLU(),\n","                                         *[conv_block(512, 512) for _ in range(3)],\n","                                         CBAM(512, r=4),\n","                                         nn.MaxPool2d(kernel_size=2, stride=2))\n","\n","        self.conv_block5 = nn.Sequential(*[conv_block(512, 512) for _ in range(4)],\n","                                         CBAM(512, r=4),\n","                                         nn.MaxPool2d(kernel_size=2, stride=2))\n","\n","        self.avg_pool = nn.AdaptiveAvgPool2d(output_size=(7,7))\n","\n","        self.linear1 = nn.Sequential(nn.Linear(in_features=7*7*512, out_features=4096, bias=True),\n","                                     nn.Dropout(0.5),\n","                                     nn.ReLU())\n","        self.linear2 = nn.Sequential(nn.Linear(in_features=4096, out_features=4096, bias=True),\n","                                     nn.Dropout(0.5),\n","                                     nn.ReLU())\n","        self.linear3 = nn.Linear(in_features=4096, out_features=self.out_channels, bias=True)\n","\n","    def forward(self, x):\n","        x = self.conv_block1(x)\n","        x = self.conv_block2(x)\n","        x = self.conv_block3(x)\n","        x = self.conv_block4(x)\n","        x = self.conv_block5(x)\n","        x = self.avg_pool(x)\n","        x = self.linear1(x.view(x.shape[0], -1))\n","        x = self.linear2(x)\n","        x = self.linear3(x)\n","        return x"]},{"cell_type":"markdown","metadata":{"id":"Ro9IAsRIE4FA"},"source":["## Train"]},{"cell_type":"code","execution_count":11,"metadata":{"executionInfo":{"elapsed":1086,"status":"ok","timestamp":1705839722461,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"triISVhHE4FA"},"outputs":[],"source":["model = VGG19_CBAM(in_channels=3, out_channels=100)\n","optimizer = optim.AdamW(model.parameters(), lr=0.0001, weight_decay=0.005)\n","criterion = nn.CrossEntropyLoss()\n","total_epochs = 50\n","device = 'cuda' if torch.cuda.is_available() else 'cpu'\n","model = model.to(device)"]},{"cell_type":"code","execution_count":12,"metadata":{"executionInfo":{"elapsed":6,"status":"ok","timestamp":1705839722461,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"g52lyNKxE4FA"},"outputs":[],"source":["def train(model, optimizer, train_loader, val_loader, criterion, total_epochs):\n","    print(\"Training Begin!\")\n","    print()\n","    best_accuracy = 0\n","\n","    for epoch in range(total_epochs):\n","        model.train()\n","\n","        for step, (images, labels) in enumerate(train_loader):\n","            images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True)\n","\n","            optimizer.zero_grad(set_to_none=True)\n","            outputs = model(images)\n","            loss = criterion(outputs, labels)\n","            loss.backward()\n","            optimizer.step()\n","\n","            if (step+1) % 30 == 0:\n","              print('Epoch: [{}/{}] | Step: [{}/{}] | Loss: {:.4f}'.format(epoch+1, total_epochs, step+1, len(train_loader), loss.item()))\n","\n","        with torch.no_grad():\n","            print(\"Validating...\")\n","            model.eval()\n","            total = 0\n","            correct = 0\n","            for (images, labels) in val_loader:\n","                images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True)\n","\n","                outputs = model(images)\n","                _, predicted = torch.max(outputs, dim=1)\n","                total += labels.size(0)\n","                correct += (predicted == labels).sum().item()\n","\n","            acc = (correct/total)*100\n","            if acc > best_accuracy:\n","                best_accuracy = acc\n","                torch.save(model.state_dict(), \"./model_cbam2.pt\")\n","            print(f\"Current Accuracy: {acc:.3f}%\")\n","            print(f\"Best Accuracy: {best_accuracy:.3f}%\")\n","            print()\n","\n","    print('Train Finished!')"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"GN-UQH6dE4FA"},"outputs":[],"source":["train(model=model, optimizer=optimizer, train_loader=train_loader, val_loader=valid_loader, criterion=criterion, total_epochs=total_epochs)"]},{"cell_type":"code","execution_count":14,"metadata":{"executionInfo":{"elapsed":298,"status":"ok","timestamp":1705840759640,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"HvFIDZN82KDC"},"outputs":[],"source":["def test(model, test_loader, criterion):\n","    print(\"Evaluating Test Dataset...\")\n","    print()\n","    best_accuracy = 0\n","\n","    with torch.no_grad():\n","      model.eval()\n","      total = 0\n","      correct = 0\n","      count = 0\n","\n","      for (images, labels) in test_loader:\n","          images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True)\n","\n","          outputs = model(images)\n","\n","          # top-1 accuracy\n","          _, predicted = torch.max(outputs, dim=1)\n","          total += labels.shape[0]\n","          correct += torch.eq(predicted, labels).sum().item()\n","\n","          # top-5 accuracy\n","          _, predicted = torch.topk(outputs, 5)\n","          for gt, pred in zip(labels, predicted):\n","            if gt in pred:\n","              count += 1\n","\n","      top1_acc = (correct/total)*100\n","      top5_acc = (count/total)*100\n","      print(f\"Test Top-1 Accuracy: {top1_acc:.2f}%\")\n","      print(f\"Test Top-5 Accuracy: {top5_acc:.2f}%\")\n"]},{"cell_type":"code","execution_count":14,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":3520,"status":"ok","timestamp":1705836787120,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"7CrmXFMt0_b6","outputId":"a8121f24-6e17-4004-fbe1-5d8ab608dbcb"},"outputs":[{"name":"stdout","output_type":"stream","text":["Evaluating Test Dataset...\n","\n","Test Top-1 Accuracy: 39.87%\n","Test Top-5 Accuracy: 67.69%\n"]}],"source":["model = VGG19(in_channels=3, out_channels=100).cuda()\n","model.load_state_dict(torch.load(\"./model.pt\"))\n","test(model=model, test_loader=test_loader, criterion=criterion) # No CBAM"]},{"cell_type":"code","execution_count":15,"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"elapsed":3637,"status":"ok","timestamp":1705839621816,"user":{"displayName":"임재규","userId":"04429980081645298585"},"user_tz":-540},"id":"N6S6pL0H_5wh","outputId":"63612b39-95d9-4b63-8e70-74ea8a0987a3"},"outputs":[{"name":"stdout","output_type":"stream","text":["Evaluating Test Dataset...\n","\n","Test Top-1 Accuracy: 40.95%\n","Test Top-5 Accuracy: 68.24%\n"]}],"source":["model = VGG19_CBAM(in_channels=3, out_channels=100).cuda()\n","model.load_state_dict(torch.load(\"./model_cbam.pt\"))\n","test(model=model, test_loader=test_loader, criterion=criterion) # With CBAM"]}],"metadata":{"accelerator":"GPU","colab":{"gpuType":"A100","provenance":[]},"kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.8.5"},"orig_nbformat":4},"nbformat":4,"nbformat_minor":0}