AnoMed_PyTorch_ResNet_MLP.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install tensorboard datasets -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torch.optim.lr_scheduler import ExponentialLR\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torchvision import transforms\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from datasets import load_dataset\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.models as models\n",
    "\n",
    "# Load pre-trained ResNet model\n",
    "resnet_model = models.resnet50(pretrained=True)\n",
    "\n",
    "# Modify the final layer to match the number of attributes (40 in this case)\n",
    "num_features = resnet_model.fc.in_features\n",
    "resnet_model.fc = nn.Linear(num_features, 40) # 2048 in, 40 out\n",
    "checkpoint = torch.load(\"./trained/model_resnet.pth\")\n",
    "\n",
    "resnet_model.load_state_dict(checkpoint)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleMLP(nn.Module):\n",
    "    def __init__(self, input_dim=40):\n",
    "        super(SimpleMLP, self).__init__()\n",
    "        \n",
    "        self.fc1 = nn.Linear(input_dim, 64)\n",
    "        self.bn1 = nn.BatchNorm1d(64)\n",
    "        self.relu1 = nn.LeakyReLU()\n",
    "        \n",
    "        self.fc2 = nn.Linear(64, 128)\n",
    "        self.bn2 = nn.BatchNorm1d(128)\n",
    "        self.relu2 = nn.LeakyReLU()\n",
    "\n",
    "        self.fc3 = nn.Linear(128, 256)\n",
    "        self.bn3 = nn.BatchNorm1d(256)\n",
    "        self.relu3 = nn.LeakyReLU()\n",
    "\n",
    "        self.fc4 = nn.Linear(256, 512)\n",
    "        self.bn4 = nn.BatchNorm1d(512)\n",
    "        self.relu4 = nn.LeakyReLU()\n",
    "\n",
    "        self.fc5 = nn.Linear(512, 1024)\n",
    "        self.bn5 = nn.BatchNorm1d(1024)\n",
    "        self.relu5 = nn.LeakyReLU()\n",
    "\n",
    "        self.fc6 = nn.Linear(1024, 1024)\n",
    "        self.bn6 = nn.BatchNorm1d(1024)\n",
    "        self.relu6 = nn.LeakyReLU()\n",
    "\n",
    "        self.fc7 = nn.Linear(1024, 1024)\n",
    "        self.bn7 = nn.BatchNorm1d(1024)\n",
    "        self.relu7 = nn.LeakyReLU()\n",
    "\n",
    "        self.fc8 = nn.Linear(1024, 128)\n",
    "        self.bn8 = nn.BatchNorm1d(128)\n",
    "        self.relu8 = nn.LeakyReLU()\n",
    "\n",
    "        self.fc9 = nn.Linear(128, 64)\n",
    "        self.bn9 = nn.BatchNorm1d(64)\n",
    "        self.relu9 = nn.LeakyReLU()\n",
    "\n",
    "        self.fc10 = nn.Linear(64, 32)\n",
    "        self.bn10 = nn.BatchNorm1d(32)\n",
    "        self.relu10 = nn.LeakyReLU()\n",
    "\n",
    "        self.fc11 = nn.Linear(32, 16)\n",
    "        self.bn11 = nn.BatchNorm1d(16)\n",
    "        self.relu11 = nn.LeakyReLU()\n",
    "\n",
    "        self.fc_out = nn.Linear(16, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # First layer block\n",
    "        x1 = self.fc1(x)\n",
    "        x1 = self.bn1(x1)\n",
    "        x1 = self.relu1(x1)\n",
    "        \n",
    "        # Second layer block\n",
    "        x2 = self.fc2(x1)\n",
    "        x2 = self.bn2(x2)\n",
    "        x2 = self.relu2(x2)\n",
    "\n",
    "        # Skip connection from x1 to x2\n",
    "        x2 = x2 + x1\n",
    "\n",
    "        # Third layer block\n",
    "        x3 = self.fc3(x2)\n",
    "        x3 = self.bn3(x3)\n",
    "        x3 = self.relu3(x3)\n",
    "\n",
    "        # Skip connection from x2 to x3\n",
    "        x3 = x3 + x2\n",
    "\n",
    "        # Continue with the remaining layers\n",
    "        x4 = self.fc4(x3)\n",
    "        x4 = self.bn4(x4)\n",
    "        x4 = self.relu4(x4)\n",
    "\n",
    "        x5 = self.fc5(x4)\n",
    "        x5 = self.bn5(x5)\n",
    "        x5 = self.relu5(x5)\n",
    "\n",
    "        # Add skip connection from x4 to x5\n",
    "        x5 = x5 + x4\n",
    "\n",
    "        x6 = self.fc6(x5)\n",
    "        x6 = self.bn6(x6)\n",
    "        x6 = self.relu6(x6)\n",
    "\n",
    "        x7 = self.fc7(x6)\n",
    "        x7 = self.bn7(x7)\n",
    "        x7 = self.relu7(x7)\n",
    "\n",
    "        x8 = self.fc8(x7)\n",
    "        x8 = self.bn8(x8)\n",
    "        x8 = self.relu8(x8)\n",
    "\n",
    "        # Skip connection from x7 to x8\n",
    "        x8 = x8 + x7\n",
    "\n",
    "        x9 = self.fc9(x8)\n",
    "        x9 = self.bn9(x9)\n",
    "        x9 = self.relu9(x9)\n",
    "\n",
    "        x10 = self.fc10(x9)\n",
    "        x10 = self.bn10(x10)\n",
    "        x10 = self.relu10(x10)\n",
    "\n",
    "        x11 = self.fc11(x10)\n",
    "        x11 = self.bn11(x11)\n",
    "        x11 = self.relu11(x11)\n",
    "\n",
    "        output = self.fc_out(x11)\n",
    "\n",
    "        return output\n",
    "\n",
    "mlp = SimpleMLP(input_dim=40)\n",
    "checkpoint = torch.load(\"./trained/model_hairline.pth\")\n",
    "mlp.load_state_dict(checkpoint)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "\n",
    "# Define custom transform\n",
    "custom_transform = transforms.Compose([\n",
    "    transforms.CenterCrop((178, 178)),\n",
    "    transforms.Resize((128, 128)),\n",
    "    transforms.ToTensor()\n",
    "])\n",
    "\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "class CelebADataset(Dataset):\n",
    "    def __init__(self, images, targets, transform=None):\n",
    "        self.images = images  # List of image paths or byte data dictionaries\n",
    "        self.targets = targets  # Corresponding targets\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.images)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        image_data = self.images[idx]\n",
    "\n",
    "        if isinstance(image_data, dict):  # If image data is stored as a dict (e.g., {'bytes': ...})\n",
    "            img_bytes = image_data['bytes']\n",
    "            image = Image.open(BytesIO(img_bytes))\n",
    "        else:  # Otherwise, assume it's a file path\n",
    "            image = Image.open(image_data)\n",
    "\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "\n",
    "        target = self.targets[idx]\n",
    "        return image, target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "def eval_model(model, val_loader, criterion, num_epochs=1):\n",
    "    writer = SummaryWriter(log_dir=\"./tensorboard_logs\")\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        model.eval()\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for batch_idx, (images, targets) in enumerate(val_loader):\n",
    "                images, targets = images.to(device), targets.to(device)\n",
    "\n",
    "                outputs = model(images)\n",
    "                loss = criterion(outputs, targets)\n",
    "                eval_loss += loss.item()\n",
    "\n",
    "                predictions = torch.sigmoid(outputs)\n",
    "                predicted_labels = (predictions > 0.5).float()\n",
    "                correct_predictions += (predicted_labels == targets).sum().item()\n",
    "                total_predictions += targets.numel()\n",
    "\n",
    "        avg_eval_loss = eval_loss / len(val_loader)\n",
    "        accuracy = correct_predictions / total_predictions\n",
    "\n",
    "        # Log to TensorBoard\n",
    "        writer.add_scalar(\"Loss/Eval\", avg_eval_loss, epoch)\n",
    "        writer.add_scalar(\"Accuracy\", accuracy, epoch)\n",
    "\n",
    "        print(f\"Epoch {epoch + 1}, Eval Loss: {avg_eval_loss:.4f}, Accuracy: {accuracy:.4f}\")\n",
    "\n",
    "    # Close the TensorBoard writer\n",
    "    writer.close()\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = \"./data\"\n",
    "celeba = load_dataset(\"tpremoli/CelebA-attrs\", split=\"all\", cache_dir=data_dir)\n",
    "\n",
    "df = celeba.to_pandas()\n",
    "\n",
    "attr_names = df.columns[1:-1]\n",
    "image_col_name = df.columns[0]\n",
    "\n",
    "attributes = df[attr_names].values.astype(np.float32)\n",
    "\n",
    "inputs = df[image_col_name].values\n",
    "\n",
    "targets = (attributes + 1) / 2\n",
    "\n",
    "train_percentage = 0.8\n",
    "val_percentage = 0.1\n",
    "test_percentage = 1 - train_percentage - val_percentage\n",
    "\n",
    "training_amount = int(train_percentage * len(inputs))\n",
    "validation_amount = int(val_percentage * len(inputs))\n",
    "testing_amount = int(train_percentage * len(inputs))\n",
    "\n",
    "train_inputs = inputs[0:training_amount]\n",
    "val_inputs = inputs[training_amount:training_amount+validation_amount]\n",
    "test_inputs = inputs[training_amount+validation_amount:]\n",
    "\n",
    "\n",
    "train_targets = targets[0:training_amount]\n",
    "val_targets = targets[training_amount:training_amount+validation_amount]\n",
    "test_targets = targets[training_amount+validation_amount:]\n",
    "\n",
    "train_targets = torch.tensor(train_targets, dtype=torch.float32).to(device)\n",
    "val_targets = torch.tensor(val_targets, dtype=torch.float32).to(device)\n",
    "test_targets = torch.tensor(test_targets, dtype=torch.float32).to(device)\n",
    "\n",
    "# Custom dataset and dataloaders for ResNet and CNN\n",
    "train_dataset = CelebADataset(train_inputs, train_targets, transform=custom_transform)\n",
    "val_dataset = CelebADataset(val_inputs, val_targets, transform=custom_transform)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CombinedModel(nn.Module):\n",
    "\n",
    "    def __init__(self, resnet_model, mlp):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.resnet = resnet_model # Last layer is from FC 40\n",
    "        self.mlp = mlp # Input expects 40\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        resnet_output = self.resnet(x)\n",
    "        output = self.mlp(resnet_output)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_model = CombinedModel(resnet_model, mlp)\n",
    "\n",
    "eval_model(combined_model, val_loader, nn.BCEWithLogitsLoss())"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}