{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[],"machine_shape":"hm","gpuType":"L4","authorship_tag":"ABX9TyOWgzzWvZ3woEEb4aaAvfnw"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"},"accelerator":"GPU"},"cells":[{"cell_type":"code","source":["from google.colab import drive\n","\n","\n","drive.mount('/content/gdrive', force_remount=True)"],"metadata":{"id":"vB5ic-v8hMJb"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import torch\n","from torch.utils.data import DataLoader, TensorDataset, random_split\n","\n","test_images = torch.load('./test_data.pt')\n","test_labels = torch.load('./test_labels.pt')\n","train_images = torch.load('./train_data.pt')\n","train_labels = torch.load('./train_labels.pt')\n","\n","\n","full_train_dataset = TensorDataset(train_images, train_labels)\n","\n","# split train\n","train_size = int(0.8 * len(full_train_dataset))\n","val_size = len(full_train_dataset) - train_size\n","train_dataset, val_dataset = random_split(full_train_dataset, [train_size, val_size])\n","test_dataset = TensorDataset(test_images, test_labels)\n","\n","\n","batch_size = 32\n","shuffle = True\n","\n","\n","train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=shuffle)\n","val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=False)\n","test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)\n","\n","\n","print(f'Train dataset size: {len(train_dataset)}')\n","print(f'Val dataset size: {len(val_dataset)}')\n","print(f'Test dataset size: {len(test_dataset)}')\n","print(f'batch size: {batch_size}')"],"metadata":{"id":"7mSZR2uFiruX"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","from torchvision import models\n","\n","class VGG16_malimg(nn.Module):\n","    def __init__(self, num_classes=512):\n","        super(VGG16_malimg, self).__init__()\n","        # load vgg16\n","        vgg16 = models.vgg16(pretrained=True)\n","        self.features = vgg16.features\n","\n","        # Define fc\n","        self.classifier = nn.Sequential(\n","            nn.Linear(512 * 7 * 7, 4096),\n","            nn.ReLU(True),\n","            nn.Dropout(),\n","            nn.Linear(4096, 4096),\n","            nn.ReLU(True),\n","            nn.Dropout(),\n","            nn.Linear(4096, num_classes),\n","        )\n","\n","    def forward(self, x, extra_input):# 56 * 56 * 128\n","\n","        x = self.features[:10](x)\n","\n","        x = x * extra_input\n","\n","\n","        x = self.features[10:](x)\n","\n","\n","        x = torch.flatten(x, 1)\n","\n","\n","        x = self.classifier(x)\n","\n","        return x\n"],"metadata":{"id":"7FWLgNl1L7Qw"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class ResNet18_malimg(nn.Module):\n","    def __init__(self, pretrained_model=None):\n","        super(ResNet18_malimg, self).__init__()\n","\n","        if pretrained_model is None:\n","            pretrained_model = models.resnet18(pretrained=True)\n","\n","        # Load ResNet\n","        self.original_model = pretrained_model\n","\n","    def forward(self, x, extra_input):#56 * 56 * 64\n","        x = self.original_model.conv1(x)\n","        x = self.original_model.bn1(x)\n","        x = self.original_model.relu(x)\n","        x = self.original_model.maxpool(x)\n","\n","        x = self.original_model.layer1(x)\n","\n","        x = x * extra_input\n","\n","        x = self.original_model.layer2(x)\n","\n","        x = self.original_model.layer3(x)\n","\n","        x = self.original_model.layer4(x)\n","\n","        x = self.original_model.avgpool(x)\n","        x = torch.flatten(x, 1)\n","\n","\n","        return x"],"metadata":{"id":"ypTIivKgNOug"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","from torchvision import models\n","\n","class EfficientNetB0_malimg(nn.Module):\n","    def __init__(self, num_classes=512):\n","        super(EfficientNetB0_malimg, self).__init__()\n","        # Load EfficientNetB0\n","        efficientnet_b0 = models.efficientnet_b0(pretrained=True)\n","        self.features = efficientnet_b0.features\n","\n","        # Define fc\n","        self.classifier = nn.Sequential(\n","            nn.Linear(1280, 512),\n","            nn.ReLU(True),\n","            nn.Dropout(),\n","            nn.Linear(512, num_classes),\n","        )\n","\n","    def forward(self, x, extra_input):# 112 * 112 * 16\n","\n","        x = self.features[:2](x)\n","\n","\n","        x = x * extra_input\n","\n","\n","        x = self.features[2:](x)\n","\n","\n","        x = F.adaptive_avg_pool2d(x, 1).reshape(x.size(0), -1)\n","\n","\n","        x = self.classifier(x)\n","\n","        return x\n"],"metadata":{"id":"QrJkRyfCOmZV"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class VGG16_opimg(nn.Module):\n","    def __init__(self, num_classes=512):\n","        super(VGG16_opimg, self).__init__()\n","        # Load VGG\n","        vgg16 = models.vgg16(pretrained=True)\n","        self.features = vgg16.features\n","\n","        # Define fc\n","        self.classifier = nn.Sequential(\n","            nn.Linear(512 * 7 * 7, 4096),\n","            nn.ReLU(True),\n","            nn.Dropout(),\n","            nn.Linear(4096, 4096),\n","            nn.ReLU(True),\n","            nn.Dropout(),\n","            nn.Linear(4096, num_classes),\n","        )\n","\n","    def forward(self, x, extra_input):# 14 * 14 * 512\n","\n","        x = self.features[:30](x)\n","\n","        features = x\n","\n","\n","        x = self.features[30:](x)\n","\n","\n","        x = torch.flatten(x, 1)\n","\n","\n","        x = self.classifier(x)\n","\n","        return x, features"],"metadata":{"id":"oQ7rcWaibUeB"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class ResNet18_opimg(nn.Module):\n","    def __init__(self, pretrained_model=None):\n","        super(ResNet18_opimg, self).__init__()\n","\n","        if pretrained_model is None:\n","            pretrained_model = models.resnet18(pretrained=True)\n","\n","\n","        self.original_model = pretrained_model\n","\n","    def forward(self, x, extra_input):#7 * 7 * 512\n","        x = self.original_model.conv1(x)\n","        x = self.original_model.bn1(x)\n","        x = self.original_model.relu(x)\n","        x = self.original_model.maxpool(x)\n","\n","        x = self.original_model.layer1(x)\n","\n","        x = self.original_model.layer2(x)\n","\n","        x = self.original_model.layer3(x)\n","\n","        x = self.original_model.layer4(x)\n","\n","        features = x\n","\n","        x = self.original_model.avgpool(x)\n","        x = torch.flatten(x, 1)\n","\n","\n","        return x, features"],"metadata":{"id":"BAg-tysYbV9S"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","from torchvision import models\n","class EfficientNetB0_opimg(nn.Module):\n","    def __init__(self):\n","        super(EfficientNetB0_opimg, self).__init__()\n","        self.original_model = models.efficientnet_b0(pretrained=True)\n","        self.dropout = nn.Dropout(p=0.2, inplace=True)\n","        self.linear = nn.Linear(in_features=1280, out_features=512, bias=True)\n","\n","    def forward(self, x):# 7 * 7 * 1280\n","\n","        features = self.original_model.features(x)\n","\n","\n","        x = self.original_model.avgpool(features)\n","        x = torch.flatten(x, 1)\n","        x = self.dropout(x)\n","        output = self.linear(x)\n","\n","\n","        return output, features"],"metadata":{"id":"85v54hQuWNiz"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","from torchvision import models\n","\n","class VGG16_original(nn.Module):\n","    def __init__(self, num_classes=512):\n","        super(VGG16_original, self).__init__()\n","\n","        vgg16 = models.vgg16(pretrained=True)\n","        self.features = vgg16.features\n","\n","\n","        self.classifier = nn.Sequential(\n","            nn.Linear(512 * 7 * 7, 4096),\n","            nn.ReLU(True),\n","            nn.Dropout(),\n","            nn.Linear(4096, 4096),\n","            nn.ReLU(True),\n","            nn.Dropout(),\n","            nn.Linear(4096, num_classes),\n","        )\n","\n","    def forward(self, x):# 56 * 56 * 128\n","\n","        x = self.features(x)\n","\n","\n","\n","        x = torch.flatten(x, 1)\n","\n","\n","        x = self.classifier(x)\n","\n","        return x\n"],"metadata":{"id":"dTkvDFdSs_Co"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class ResNet18_original(nn.Module):\n","    def __init__(self, pretrained_model=None):\n","        super(ResNet18_original, self).__init__()\n","\n","        if pretrained_model is None:\n","            pretrained_model = models.resnet18(pretrained=True)\n","\n","\n","        self.original_model = pretrained_model\n","\n","    def forward(self, x):#56 * 56 * 64\n","        x = self.original_model.conv1(x)\n","        x = self.original_model.bn1(x)\n","        x = self.original_model.relu(x)\n","        x = self.original_model.maxpool(x)\n","\n","        x = self.original_model.layer1(x)\n","\n","        x = self.original_model.layer2(x)\n","\n","        x = self.original_model.layer3(x)\n","\n","        x = self.original_model.layer4(x)\n","\n","        x = self.original_model.avgpool(x)\n","        x = torch.flatten(x, 1)\n","\n","\n","        return x"],"metadata":{"id":"5Q1cdfPatQZ-"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","from torchvision import models\n","\n","class EfficientNetB0_original(nn.Module):\n","    def __init__(self, num_classes=512):\n","        super(EfficientNetB0_original, self).__init__()\n","\n","        efficientnet_b0 = models.efficientnet_b0(pretrained=True)\n","        self.features = efficientnet_b0.features\n","\n","\n","        self.classifier = nn.Sequential(\n","            nn.Linear(1280, 512),\n","            nn.ReLU(True),\n","            nn.Dropout(),\n","            nn.Linear(512, num_classes),\n","        )\n","\n","    def forward(self, x):\n","\n","        x = self.features[:2](x)\n","\n","\n","        x = self.features[2:](x)\n","\n","\n","        x = F.adaptive_avg_pool2d(x, 1).reshape(x.size(0), -1)\n","\n","\n","        x = self.classifier(x)\n","\n","        return x\n"],"metadata":{"id":"o9AMbTY3tanE"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class CrossAttention(nn.Module):\n","    def __init__(self, embed_size=1):\n","        super(CrossAttention, self).__init__()\n","        self.embed_size = embed_size\n","\n","        self.values1 = nn.Linear(embed_size, embed_size, bias=False)\n","        self.keys1 = nn.Linear(embed_size, embed_size, bias=False)\n","        self.queries1 = nn.Linear(embed_size, embed_size, bias=False)\n","        #self.fc_out = nn.Linear(embed_size, embed_size)\n","\n","        self.values2 = nn.Linear(embed_size, embed_size, bias=False)\n","        self.keys2 = nn.Linear(embed_size, embed_size, bias=False)\n","        self.queries2 = nn.Linear(embed_size, embed_size, bias=False)\n","\n","    def forward(self, values1, keys1, query1, values2, keys2, query2, mask=None):\n","        N = query1.shape[0]\n","\n","        L1 = values1.shape[1]\n","        L2 = values2.shape[1]\n","        D = self.embed_size\n","\n","        values1 = self.values1(values1)\n","        keys1 = self.keys1(keys1)\n","        queries1 = self.queries1(query1)\n","\n","        values2 = self.values2(values2)\n","        keys2 = self.keys2(keys2)\n","        queries2 = self.queries2(query2)\n","\n","\n","        energy1 = torch.bmm(queries1, keys2.transpose(1, 2)) / (self.embed_size ** (1 / 2))\n","\n","        energy2 = torch.bmm(queries2, keys1.transpose(1, 2)) / (self.embed_size ** (1 / 2))\n","\n","\n","\n","        attention1 = torch.softmax(energy1, dim=2)\n","        attention2 = torch.softmax(energy2, dim=2)\n","\n","        out1 = torch.bmm(attention1, values1)\n","        out2 = torch.bmm(attention2, values2)\n","\n","\n","\n","\n","\n","        return out1, out2"],"metadata":{"id":"g2TE1TfSsXF-"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import torch\n","import torch.nn as nn\n","\n","class MultilayerDeepNetwork(nn.Module):\n","    def __init__(self):\n","        super(MultilayerDeepNetwork, self).__init__()\n","\n","        self.fc1 = nn.Linear(512, 256)\n","        self.fc2 = nn.Linear(256, 256)\n","        self.fc3 = nn.Linear(1024, 512)\n","        self.fc4 = nn.Linear(1024, 512)\n","        self.fc5 = nn.Linear(1024, 512)\n","        self.fc6 = nn.Linear(512, 256)\n","        self.fc7 = nn.Linear(256, 256)\n","\n","    def forward(self, x1, x2):\n","\n","        layer1_output1 = torch.relu(self.fc1(x1))\n","        layer1_output2 = torch.relu(self.fc3(torch.cat((x1, x2), dim=1)))\n","        layer1_output3 = torch.relu(self.fc6(x2))\n","\n","        layer2_output1 = torch.relu(self.fc2(layer1_output1))\n","        layer2_output2 = torch.relu(self.fc4(torch.cat((layer1_output1, layer1_output2, layer1_output3), dim=1)))\n","        layer2_output3 = torch.relu(self.fc7(layer1_output3))\n","\n","        layer3_output = torch.relu(self.fc5(torch.cat((layer2_output1, layer2_output2, layer2_output3), dim=1)))\n","\n","        return layer3_output"],"metadata":{"id":"ZY2MUjRnLu4Q"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import torch\n","import torch.nn as nn\n","import torchvision.models as models\n","import torch.optim as optim\n","\n","\n","\n","class MultilayerFeatureFusionModel(nn.Module):\n","    def __init__(self, num_classes=9):\n","        super(MultilayerFeatureFusionModel, self).__init__()\n","        self.malimg_extractor = ResNet18_malimg()\n","        self.opimg_extractor = EfficientNetB0_opimg()\n","        self.cross_attention = CrossAttention(embed_size=1)\n","        self.fc = nn.Linear(512 * 2, num_classes)\n","        self.norm = nn.LayerNorm(512)\n","        self.dropout = nn.Dropout(0.3)\n","        self.conv_transpose = nn.ConvTranspose2d(\n","    in_channels=1280,\n","    out_channels=64,\n","    kernel_size=8,\n","    stride=8,\n","    padding=0,\n","    output_padding=0,\n","    bias=True\n",")\n","\n","\n","    def forward(self, x):\n","\n","        x1 = x[:, :3, :, :]\n","\n","        x2 = x[:, 3:, :, :]\n","\n","\n","\n","        x2, extra_features = self.opimg_extractor(x2)\n","\n","        extra_features = self.conv_transpose(extra_features)\n","\n","        x1 = self.malimg_extractor(x1, extra_features)\n","\n","        value1 = x1.unsqueeze(2)\n","        value2 = x2.unsqueeze(2)\n","\n","        # Cross Attention\n","        out1, out2 = self.cross_attention(value1, value1, value1, value2, value2, value2)\n","\n","        # Reshape the outputs\n","        out1 = out1.reshape(out1.size(0), out1.size(1) * out1.size(2))\n","        out2 = out2.reshape(out2.size(0), out2.size(1) * out2.size(2))\n","\n","        x1 = self.norm(x1 + out1)\n","        x2 = self.norm(x2 + out2)\n","\n","        x1 = self.dropout(x1)\n","        x2 = self.dropout(x2)\n","\n","\n","        x = torch.cat((x1, x2), dim=1)\n","\n","\n","        out = self.fc(x)\n","        return out\n","\n"],"metadata":{"id":"nRHsfu8knk_U"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import torch\n","import torch.nn as nn\n","import torch.optim as optim\n","from torch.utils.data import DataLoader\n","from sklearn.metrics import classification_report, confusion_matrix, precision_score, recall_score, f1_score\n","import matplotlib.pyplot as plt\n","import seaborn as sns\n","import numpy as np\n","\n","# Define Device\n","device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n","\n","\n","model = MultilayerFeatureFusionModel()\n","model = model.to(device)\n","\n","# Define loss and optimizer\n","criterion = nn.CrossEntropyLoss()\n","optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)\n","\n","num_epochs = 1\n","best_model_wts = model.state_dict()\n","best_val_acc = 0.0\n","\n","for epoch in range(num_epochs):\n","    print(f'Epoch {epoch}/{num_epochs - 1}')\n","    print('-' * 10)\n","\n","    # Train\n","    model.train()\n","    running_loss = 0.0\n","    running_corrects = 0\n","\n","    for inputs, labels in train_loader:\n","        inputs = inputs.to(device)\n","        labels = labels.to(device)\n","\n","        optimizer.zero_grad()\n","\n","        with torch.set_grad_enabled(True):\n","            outputs = model(inputs)\n","            _, preds = torch.max(outputs, 1)\n","            loss = criterion(outputs, labels)\n","\n","            loss.backward()\n","            optimizer.step()\n","\n","        running_loss += loss.item() * inputs.size(0)\n","        running_corrects += torch.sum(preds == labels.data)\n","\n","    epoch_loss = running_loss / len(train_loader.dataset)\n","    epoch_acc = running_corrects.double() / len(train_loader.dataset)\n","\n","    print(f'train Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')\n","\n","    # evaluate\n","    model.eval()\n","    running_loss = 0.0\n","    running_corrects = 0\n","    all_preds = []\n","    all_labels = []\n","\n","    for inputs, labels in val_loader:\n","        inputs = inputs.to(device)\n","        labels = labels.to(device)\n","\n","        with torch.set_grad_enabled(False):\n","            outputs = model(inputs)\n","            _, preds = torch.max(outputs, 1)\n","            loss = criterion(outputs, labels)\n","\n","        running_loss += loss.item() * inputs.size(0)\n","        running_corrects += torch.sum(preds == labels.data)\n","\n","        all_preds.extend(preds.cpu().numpy())\n","        all_labels.extend(labels.cpu().numpy())\n","\n","    epoch_loss = running_loss / len(val_loader.dataset)\n","    epoch_acc = running_corrects.double() / len(val_loader.dataset)\n","    epoch_precision = precision_score(all_labels, all_preds, average='weighted')\n","    epoch_recall = recall_score(all_labels, all_preds, average='weighted')\n","    epoch_f1 = f1_score(all_labels, all_preds, average='weighted')\n","\n","    print(f'val Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f} Precision: {epoch_precision:.4f} Recall: {epoch_recall:.4f} F1: {epoch_f1:.4f}')\n","\n","\n","    if epoch_acc > best_val_acc:\n","        best_val_acc = epoch_acc\n","        best_model_wts = model.state_dict()\n","\n","\n","model.load_state_dict(best_model_wts)\n","\n","# Test\n","model.eval()\n","running_loss = 0.0\n","running_corrects = 0\n","all_preds = []\n","all_labels = []\n","\n","for inputs, labels in test_loader:\n","    inputs = inputs.to(device)\n","    labels = labels.to(device)\n","\n","    with torch.set_grad_enabled(False):\n","        outputs = model(inputs)\n","        _, preds = torch.max(outputs, 1)\n","        loss = criterion(outputs, labels)\n","\n","    running_loss += loss.item() * inputs.size(0)\n","    running_corrects += torch.sum(preds == labels.data)\n","\n","    all_preds.extend(preds.cpu().numpy())\n","    all_labels.extend(labels.cpu().numpy())\n","\n","epoch_loss = running_loss / len(test_loader.dataset)\n","epoch_acc = running_corrects.double() / len(test_loader.dataset)\n","epoch_precision = precision_score(all_labels, all_preds, average='weighted')\n","epoch_recall = recall_score(all_labels, all_preds, average='weighted')\n","epoch_f1 = f1_score(all_labels, all_preds, average='weighted')\n","\n","print(f'test Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f} Precision: {epoch_precision:.4f} Recall: {epoch_recall:.4f} F1: {epoch_f1:.4f}')\n","\n","\n","conf_matrix = confusion_matrix(all_labels, all_preds)\n","plt.figure(figsize=(10, 8))\n","sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues', cbar=True)\n","plt.title('Confusion Matrix')\n","plt.xlabel('Predicted Label')\n","plt.ylabel('True Label')\n","plt.show()\n","\n","\n","torch.save(model, './best_model')\n","print(f'Best val Acc: {best_val_acc:.4f}')\n"],"metadata":{"id":"RW4FyhBdp8Pu"},"execution_count":null,"outputs":[]}]}