{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":3004,"databundleVersionId":861823,"sourceType":"competition"},{"sourceId":1512654,"sourceType":"datasetVersion","datasetId":891340}],"dockerImageVersionId":30665,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.autograd import Variable\n\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader, Dataset\nimport torchvision.transforms as transforms\n\nimport os\nfrom PIL import Image\n\nimport numpy as np\nfrom tqdm.notebook import tqdm\nimport matplotlib.pyplot as plt\n\n\nclass ImageLoader(Dataset):\n    def __init__(self, split, transform=None):\n\n        # Update these paths\n        healthy_path = f\"/kaggle/input/cotton-disease-dataset/Cotton Disease/{split}/fresh cotton leaf\"\n        diseased_path = f\"/kaggle/input/cotton-disease-dataset/Cotton Disease/{split}/diseased cotton leaf\"\n\n        healthy_files = [os.path.join(healthy_path, file)\n                         for file in os.listdir(healthy_path)]\n        diseased_files = [os.path.join(diseased_path, file)\n                          for file in os.listdir(diseased_path)]\n\n        self.data = [(file, 0.0) for file in healthy_files] + \\\n            [(file, 1.0) for file in diseased_files]\n\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        image_path, label = self.data[idx]\n        image = (Image.open(image_path).convert(\"RGB\"))\n        \n        image = transforms.ToTensor()(image)\n        \n        image -= torch.min(image)\n        image /= torch.max(image)\n        \n        if self.transform:\n            image = self.transform(image)\n\n        return image, torch.tensor([1.0 - label, label])\n\n\ndef squash(x, axis=-1):\n    norm = torch.norm(x, p=2, dim=axis, keepdim=True)\n    scale = norm**2 / (1 + norm**2) / (norm + 1e-8)\n    return scale * x\n\n\nclass PrimaryCapsule(nn.Module):\n    \"\"\"\n    Apply Conv2D with `out_channels` and then reshape to get capsules\n    :param in_channels: input channels\n    :param out_channels: output channels\n    :param dim_caps: dimension of capsule\n    :param kernel_size: kernel size\n    :return: output tensor, size=[batch, num_caps, dim_caps]\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels, dim_caps, kernel_size, stride=1, padding=0):\n        super().__init__()\n        self.dim_caps = dim_caps\n        self.conv2d = nn.Conv2d(\n            in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding)\n\n    def forward(self, x):\n        outputs = self.conv2d(x)\n        outputs = outputs.view(x.size(0), -1, self.dim_caps)\n        return squash(outputs)\n\n\nclass DenseCapsule(nn.Module):\n    \"\"\"\n    The dense capsule layer. It is similar to Dense (FC) layer. Dense layer has `in_num` inputs, each is a scalar, the output of the neuron from the former layer, and it has `out_num` output neurons. DenseCapsule just expands the output of the neuron from scalar to vector. So its input size  = [None, in_num_caps  , in_dim_caps] and \n    output size = [None, out_num_caps, out_dim_caps]. For Dense Layer, in_dim_caps = out_dim_caps = 1.\n\n    :param in_num_caps: number of capsules inputted to this layer\n    :param in_dim_caps: dimension of input capsules\n    :param out_num_caps: number of capsules outputted from this layer\n    :param out_dim_caps: dimension of output capsules\n    :param routings: number of iterations for the routing algorithm\n    \"\"\"\n\n    def __init__(self, in_num_caps, in_dim_caps, out_num_caps, out_dim_caps, routings=3, device=\"cpu\"):\n        super().__init__()\n        self.in_num_caps = in_num_caps\n        self.in_dim_caps = in_dim_caps\n        self.out_num_caps = out_num_caps\n        self.out_dim_caps = out_dim_caps\n        self.routings = routings\n        self.weight = nn.Parameter(\n            0.01 * torch.randn(out_num_caps, in_num_caps, out_dim_caps, in_dim_caps))\n\n        self.DEVICE = device\n\n    def forward(self, x):\n        # x.size =    [batch, in_num_caps, in_dim_caps]\n        # expanded to [batch, 1          , in_num_caps, in_dim_caps,  1]\n        # weight.size = [out_num_caps, in_num_caps, out_dim_caps, in_dim_caps]\n        # torch.matmul: [out_dim_caps, in_dim_caps] x [in_dim_caps, 1] -> [out_dim_caps, 1]\n        # => x_hat.size = [batch, out_num_caps, in_num_caps, out_dim_caps]\n        x_hat = torch.squeeze(torch.matmul(\n            self.weight, x[:, None, :, :, None]), dim=-1)\n\n        # In forward pass, `x_hat_detached` = `x_hat`;\n        # In backward, no gradient can flow from `x_hat_detached` back to `x_hat`.\n        x_hat_detached = x_hat.detach()\n\n        # The prior for coupling coefficient, initialized as zeros.\n        # b.size = [batch, out_num_caps, in_num_caps]\n        b = Variable(torch.zeros(\n            x.size(0), self.out_num_caps, self.in_num_caps)).to(self.DEVICE)\n\n        for i in range(self.routings):\n            # c.size = [batch, out_num_caps, in_num_caps]\n            c = F.softmax(b, dim=1)\n\n            # At last iteration, use `x_hat` to compute `outputs` in order to backpropagate gradient\n            if i == self.routings - 1:\n                # c.size expanded to [batch, out_num_caps, in_num_caps, 1           ]\n                # x_hat.size     =   [batch, out_num_caps, in_num_caps, out_dim_caps]\n                # => outputs.size=   [batch, out_num_caps, 1          , out_dim_caps]\n                outputs = squash(\n                    torch.sum(c[:, :, :, None] * x_hat, dim=-2, keepdim=True))\n                # outputs = squash(torch.matmul(c[:, :, None, :], x_hat))  # alternative way\n\n            else:  # Otherwise, use `x_hat_detached` to update `b`. No gradients flow on this path.\n                outputs = squash(\n                    torch.sum(c[:, :, :, None] * x_hat_detached, dim=-2, keepdim=True))\n                # outputs = squash(torch.matmul(c[:, :, None, :], x_hat_detached))  # alternative way\n\n                # outputs.size        = [batch, out_num_caps, 1,           out_dim_caps]\n                # x_hat_detached.size = [batch, out_num_caps, in_num_caps, out_dim_caps]\n                # => b.size           = [batch, out_num_caps, in_num_caps]\n                b = b + torch.sum(outputs * x_hat_detached, dim=-1)\n\n        return torch.squeeze(outputs, dim=-2)\n\n\nclass CapsuleNet(nn.Module):\n    \"\"\"\n    A Capsule Network on MNIST.\n    :param input_size: data size = [channels, width, height]\n    :param classes: number of classes\n    :param routings: number of routing iterations\n    Shape:\n        - Input: (batch, channels, width, height), optional (batch, classes) .\n        - Output:((batch, classes), (batch, channels, width, height))\n    \"\"\"\n\n    def __init__(self, input_size, classes, routings, device=\"cpu\"):\n        super(CapsuleNet, self).__init__()\n        self.input_size = input_size\n        self.classes = classes\n        self.routings = routings\n        self.DEVICE = device\n\n        # Regular Conv Layer\n        self.init_conv = nn.Conv2d(input_size[0], 256, kernel_size=9, stride=1)\n\n        # Conv2D layer with `squash` activation, then reshape to [None, num_caps, dim_caps]\n        self.primarycaps = PrimaryCapsule(256, 256, 8, kernel_size=9, stride=2)\n\n        # Layer 3: Capsule layer. Routing algorithm works here.\n        self.digitcaps = DenseCapsule(in_num_caps=32*104*104, in_dim_caps=8,\n                                      out_num_caps=classes, out_dim_caps=16, routings=routings, device=device)\n\n        # Decoder network.\n\n        # For MNIST\n        # self.decoder = nn.Sequential(\n        #     nn.Linear(16*classes, 512),\n        #     nn.ReLU(inplace=True),\n        #     nn.Linear(512, 1024),\n        #     nn.ReLU(inplace=True),\n        #     nn.Linear(1024, input_size[0] * input_size[1] * input_size[2]),\n        #     nn.Sigmoid()\n        # )\n\n        self.decoder = nn.Sequential(\n            nn.ConvTranspose2d(in_channels=2, out_channels=16,\n                               kernel_size=3, stride=2),\n            nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(\n                in_channels=16, out_channels=32, kernel_size=4, stride=3),\n            nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(in_channels=32, out_channels=16,\n                               kernel_size=4, stride=2, padding=1),\n            nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(in_channels=16, out_channels=8,\n                               kernel_size=4, stride=2, padding=1),\n            nn.ReLU(inplace=True),\n            nn.ConvTranspose2d(in_channels=8, out_channels=3,\n                               kernel_size=4, stride=2, padding=1),\n            nn.Sigmoid()\n        )\n\n        self.relu = nn.ReLU()\n\n    def forward(self, x, y=None):\n        x = self.relu(self.init_conv(x))\n        x = self.primarycaps(x)\n        x = self.digitcaps(x)\n\n        length = x.norm(dim=-1)\n\n        if y is None:  # during testing, no label given. create one-hot coding using `length`\n            index = length.max(dim=1)[1]\n            y = Variable(torch.zeros(length.size()).scatter_(\n                1, index.view(-1, 1).cpu().data, 1.)).to(self.DEVICE)\n\n        recon_input = (x * y[:, :, None])\n        recon_input = recon_input.reshape(recon_input.size(0), 2, 4, 4)\n\n        reconstruction = self.decoder(recon_input)\n\n        return length, reconstruction\n\n\nclass CapsuleLoss(nn.Module):\n\n    \"\"\"\n    Capsule loss = Margin loss + lam_recon * reconstruction loss.\n    :param y_true: true labels, one-hot coding, size=[batch, classes]\n    :param y_pred: predicted labels by CapsNet, size=[batch, classes]\n    :param x: input data, size=[batch, channels, width, height]\n    :param x_recon: reconstructed data, size is same as `x`\n    :param lam_recon: coefficient for reconstruction loss\n    :return: Variable contains a scalar loss value.\n    \"\"\"\n\n    def __init__(self, lam_recon):\n        super().__init__()\n        self.lam_recon = lam_recon\n        self.recon_loss = nn.MSELoss()\n        self.class_loss = nn.CrossEntropyLoss()\n\n    def forward(self, y_true, y_pred, x, x_recon):\n\n        L = y_true * torch.clamp(0.9 - y_pred, min=0.) ** 2 + \\\n            0.5 * (1 - y_true) * torch.clamp(y_pred - 0.1, min=0.) ** 2\n        L_margin = L.sum(dim=1).mean()\n\n#         L_margin = self.class_loss(y_pred, y_true)\n\n        L_recon = self.recon_loss(x_recon, x)\n\n        return L_margin + self.lam_recon * L_recon\n\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint()\nprint(f\"Device has been set to: {DEVICE}\")\n\nIMAGE_SIZE = 224\nNUM_EPOCHS = 1\nNUM_WORKERS = 0\nBATCH_SIZE = 2\nLR = 0.001\nLR_DECAY = 0.9\nLAM_RECON = 0.0005 * (IMAGE_SIZE*IMAGE_SIZE)\nROUTINGS = 4\n\nprint()\nprint(\"Initialized constants.\")\n\ntransform = transforms.Compose([\n    transforms.Resize((224, 224)),\n])\n\ntrainset = ImageLoader(split=\"train\", transform=transform)\nvalset = ImageLoader(split=\"val\", transform=transform)\n\nprint()\nprint(\"Initialized datasets.\")\n\ntrainloader = DataLoader(trainset, batch_size=BATCH_SIZE,\n                         shuffle=True, num_workers=NUM_WORKERS)\nvalloader = DataLoader(valset, batch_size=BATCH_SIZE,\n                       shuffle=True, num_workers=NUM_WORKERS)\n\nprint()\nprint(\"Initialized dataloaders.\")\n\nmodel = CapsuleNet(input_size=[3, 224, 224],\n                   classes=2, routings=ROUTINGS, device=DEVICE).to(DEVICE)\ncriterion = CapsuleLoss(lam_recon=LAM_RECON).to(DEVICE)\noptimizer = optim.Adam(model.parameters(), lr=LR)\nscheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=LR_DECAY)\n\nprint()\nprint(\"Initialized model.\")\n\ntrain_loss_log = []\n\nbest_val_acc = 0.0\nbest_val_loss = None\nval_loss_log = []\n\nval_acc_log = []\n\nfor epoch in range(1, NUM_EPOCHS+1):\n\n    print()\n    print(f\"Epoch #{epoch}:\")\n\n    model.train()\n    train_loss = 0.0\n\n    for idx, (image, label) in enumerate(tqdm(trainloader)):\n        image = image.to(DEVICE)\n        label = label.to(DEVICE)\n\n        optimizer.zero_grad()\n        pred, recon = model(image, label)\n\n#         plt.figure(figsize=(20, 15))\n#         plt.subplot(1, 2, 1)\n#         plt.imshow(recon[0].detach().numpy().transpose(1, 2, 0))\n#         plt.title(\"Reconstruction\")\n#         plt.colorbar()\n#         plt.subplot(1, 2, 2)\n#         plt.imshow(image[0].detach().numpy().transpose(1, 2, 0))\n#         plt.title(\"Image\")\n#         plt.colorbar()\n#         plt.savefig(f\"./temp/reconstructed\")\n#         plt.close()\n\n        loss = criterion(label, pred, image, recon)\n        loss.backward()\n        train_loss += loss.data * image.size(0)\n        optimizer.step()\n\n    train_loss_log.append(train_loss.item() / len(trainloader.dataset))\n    print(f\"Train loss at epoch #{epoch}: {train_loss_log[-1]}\")\n    \n    print(pred.data)\n    print(label)\n\n    scheduler.step()\n\n    print()\n\n    model.eval()\n    val_loss = 0.0\n    correct_pred = 0\n\n    for idx, (image, label) in enumerate(tqdm(valloader)):\n        image = image.to(DEVICE)\n        label = label.to(DEVICE)\n\n        pred, recon = model(image)\n        val_loss += criterion(label, pred, image, recon).data * image.size(0)\n\n        pred = pred.data.max(1)[1]\n        label = label.data.max(1)[1]\n        correct_pred += pred.eq(label).cpu().sum()\n\n    val_loss_log.append(val_loss.item() / len(valloader.dataset))\n    accuracy = correct_pred / len(valloader.dataset)\n    val_acc_log.append(accuracy)\n\n    print(f\"Val loss at epoch #{epoch}: {val_loss_log[-1]}\")\n    print(f\"Val accuracy at epoch #{epoch}: {accuracy}\")\n    \n    print(pred.data)\n    print(label)\n    \n    np.save(f\"/kaggle/working/loss_logs.npy\", [train_loss_log, val_loss_log, val_acc_log])\n    \n    if best_val_loss is None or best_val_loss > val_loss_log[-1]:\n        best_val_loss = val_loss_log[-1]\n        torch.save(model.state_dict(), f\"/kaggle/working/best_loss_model_so_far_{epoch}.pth\")\n\n        print()\n        print(\"Model loss checkpointed.\")\n        print()\n\n    if accuracy > best_val_acc:\n        best_val_acc = accuracy\n        torch.save(model.state_dict(), f\"/kaggle/working/best_acc_model_so_far_{epoch}.pth\")\n\n        print()\n        print(\"Model acc checkpointed.\")\n        print()\n\n    if epoch % 10 == 0:\n        torch.save(model.state_dict(), f\"/kaggle/working/model_at_epoch_{epoch}.pth\")\n\nprint()\nprint(\"Script executed.\")","metadata":{"_uuid":"978dd353-19bf-4b2c-b802-cbdc0cccc799","_cell_guid":"7c137663-4b45-4621-9407-84dd4cb90bec","jupyter":{"outputs_hidden":false},"collapsed":false,"execution":{"iopub.status.busy":"2024-03-04T19:50:45.212518Z","iopub.execute_input":"2024-03-04T19:50:45.212850Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# !rm -rf /kaggle/working/*","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\n\nloss_logs = np.load(\"/kaggle/working/loss_logs.npy\")\n\nprint(loss_logs.shape)\n\ntrain_loss = loss_logs[0]\nval_loss = loss_logs[1]\nval_acc = loss_logs[2]\n\nplt.figure(figsize=(40, 10))\nplt.plot(train_loss)\nplt.plot(val_loss)\nplt.grid(True)\n\nplt.show()\nplt.close()\n\nprint(f\"Best val_loss at epoch #{np.argmin(loss_logs[1, :]) + 1}\")\nprint(f\"Best val_acc at epoch #{np.argmin(loss_logs[2, :]) + 1}\")\n\nfor idx, (train_loss, val_loss) in enumerate(zip(loss_logs[0, :], loss_logs[1, :], loss_logs[2, :])):\n    print(f\"Epoch #{idx+1}:\\t{train_loss}\\t{val_loss}\\t{val_acc}\")\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}