Decoding_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision.models as models\n",
    "import torch.nn.functional as F\n",
    "from torch import nn\n",
    "import math\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import cv2\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.0.1+cu117\n"
     ]
    }
   ],
   "source": [
    "print(torch.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/aandre/fmri_project\n"
     ]
    }
   ],
   "source": [
    "#to get the current working directory\n",
    "directory = os.getcwd()\n",
    "\n",
    "print(directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Decoder2(\n",
      "  (fc_output): Linear(in_features=13459, out_features=25088, bias=True)\n",
      "  (conv1): Conv2d(128, 64, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
      "  (relu1): ReLU()\n",
      "  (upsample1): UpsamplingNearest2d(scale_factor=2.0, mode='nearest')\n",
      "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (conv2): Conv2d(64, 48, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
      "  (relu2): ReLU()\n",
      "  (upsample2): UpsamplingNearest2d(scale_factor=2.0, mode='nearest')\n",
      "  (bn2): BatchNorm2d(48, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (conv3): Conv2d(48, 48, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
      "  (relu3): ReLU()\n",
      "  (upsample3): UpsamplingNearest2d(scale_factor=2.0, mode='nearest')\n",
      "  (bn3): BatchNorm2d(48, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (conv4): Conv2d(48, 3, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
      "  (sigmoid): Sigmoid()\n",
      "  (upsample4): UpsamplingNearest2d(scale_factor=2.0, mode='nearest')\n",
      "  (bn4): BatchNorm2d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "from models import decoder2\n",
    "\n",
    "model = decoder2.Decoder2()\n",
    "model.cuda()\n",
    "print(model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_input = torch.rand(13459).cuda()\n",
    "output_test = model(test_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 112, 112])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_test.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ">Dataloader decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3599\n"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "class MyDataset(Dataset):\n",
    "    def __init__(self, image_dir, label_dir, num_images_per_sample=32):\n",
    "        self.image_dir = image_dir\n",
    "        self.label_dir = label_dir\n",
    "        self.sample_ids = os.listdir(self.label_dir)\n",
    "        random.shuffle(self.sample_ids)\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.sample_ids)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        subfolder_id = self.sample_ids[idx]\n",
    "        subfolder_id = subfolder_id.replace(\"label\", \"\").split(\".\")[0]\n",
    "    \n",
    "        image_name = 'frame_{:03d}_{:03d}.jpg'.format(int(subfolder_id), 16)\n",
    "        image_path = os.path.join(self.image_dir, image_name)\n",
    "        im = cv2.imread(image_path)\n",
    "        im_resized = cv2.resize(im, (112, 112), interpolation=cv2.INTER_LINEAR)\n",
    "        im_tensor = torch.from_numpy(im_resized)\n",
    "        fmri_name = 'label{}.npy'.format(int(subfolder_id))\n",
    "        fmri_path = os.path.join(self.label_dir, fmri_name)\n",
    "        fmri = np.load(fmri_path)\n",
    "        return fmri, im_tensor\n",
    "        \n",
    "        return image_batch, label\n",
    "    \n",
    "    \n",
    "image_dir = '/media/miplab-nas2/Data3/andre/stimuli_final'\n",
    "label_dirs = 'labels2/1'\n",
    "\n",
    "train_dataset = MyDataset(image_dir, label_dirs)\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True)\n",
    "\n",
    "print(len(train_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lim tensor(2.5435, device='cuda:0', grad_fn=<MseLossBackward0>)\n",
      "tv tensor(2.2516, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "lasso tensor(92.2395, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor(20.1700, device='cuda:0', grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torchvision.models import vgg16\n",
    "from models.p3d_model import *\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "\n",
    "\n",
    "# Example data\n",
    "x = torch.randn(1, 3, 112, 112).cuda()  # Ground truth frame\n",
    "x_hat = torch.randn(1, 3, 112, 112).cuda().requires_grad_(True)  # Reconstructed frame\n",
    "\n",
    "# Example VGG model\n",
    "vgg_model = vgg16(pretrained=True)  # Select a subset of VGG layers\n",
    "vgg_model = vgg_model.cuda()\n",
    "vgg_model.eval()\n",
    "\n",
    "\n",
    "# Weight coefficients\n",
    "beta = 0.5\n",
    "gamma = 0.3\n",
    "delta = 0.2\n",
    "\n",
    "# Perceptual Similarity Loss (Lim)\n",
    "def perceptual_similarity_loss(x, x_hat, vgg_model):\n",
    "    x = x.view(1, 3, 112, 112)\n",
    "    x = x.cuda().float()  # Convert input tensor to float\n",
    "    x_hat = x_hat.view(1, 3, 112, 112)\n",
    "    x_hat = x_hat.cuda().float()  # Convert reconstructed tensor to float\n",
    "    \n",
    "    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "\n",
    "    # Normalize the tensors\n",
    "    x = normalize(x)\n",
    "    x_hat = normalize(x_hat)\n",
    "    \n",
    "    x_features = vgg_model(x)\n",
    "    x_hat_features = vgg_model(x_hat)\n",
    "    lim_loss = nn.MSELoss()(x_features, x_hat_features)\n",
    "    return lim_loss\n",
    "# Regularization Term (LR)\n",
    "\n",
    "def total_variation_loss(x_hat):\n",
    "    diff_h = torch.abs(x_hat[:, :, :, 1:] - x_hat[:, :, :, :-1])\n",
    "    diff_v = torch.abs(x_hat[:, :, 1:, :] - x_hat[:, :, :-1, :])\n",
    "    tv_loss = torch.mean(diff_h) + torch.mean(diff_v)\n",
    "    tv_loss.retain_grad()\n",
    "    return tv_loss\n",
    "\n",
    "def spatial_group_lasso_regularization(model):\n",
    "    lasso_loss = 0\n",
    "    for name, param in model.named_parameters():\n",
    "        if 'fc_output' in name:  # Assuming the first FC layer is named 'fc1'\n",
    "            lasso_loss += torch.norm(param, p=2)\n",
    "    return lasso_loss\n",
    "\n",
    "# Supervised Decoder Loss\n",
    "def supervised_decoder_loss(x, x_hat, vgg_model, beta, gamma, delta, model):\n",
    "    lim_loss = perceptual_similarity_loss(x, x_hat, vgg_model)\n",
    "    tv_loss = total_variation_loss(x_hat)\n",
    "    lasso_loss = spatial_group_lasso_regularization(model)\n",
    "    loss = beta * lim_loss + delta * (tv_loss + lasso_loss)\n",
    "    return loss\n",
    "\n",
    "supervised_decoder_loss(x, x_hat, vgg_model, beta, gamma, delta, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_hat = torch.randn(1, 3, 112, 112).cuda().requires_grad_(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.8036, device='cuda:0', grad_fn=<MseLossBackward0>)"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "perceptual_similarity_loss(x, x_hat, vgg_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.2532, device='cuda:0', grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total_variation_loss(x_hat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(92.2348, device='cuda:0', grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spatial_group_lasso_regularization(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "import math\n",
    "\n",
    "initial_lr = 1e-3\n",
    "lr_factor = 5\n",
    "lr_step_size = 25\n",
    "\n",
    "def get_learning_rate(epoch):\n",
    "    return initial_lr / math.pow(lr_factor, epoch // lr_step_size)\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr=get_learning_rate(0))\n",
    "scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=get_learning_rate)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 100\n",
    "print_interval = 10\n",
    "losses_tot = []\n",
    "min_loss = 10000000000000 #max\n",
    "\n",
    "#scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.85)\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    for i, (fmri, image) in enumerate(train_dataloader):\n",
    "        optimizer.zero_grad()\n",
    "        images = image.to(device).float()\n",
    "        fmri = fmri.to(device).float()\n",
    "        x_hat = model(fmri)\n",
    "        loss = supervised_decoder_loss(image, x_hat, vgg_model, beta, gamma, delta)\n",
    "        losses_tot.append(loss.mean().item())\n",
    "        if (loss.mean().item()) < min_loss:\n",
    "            min_loss = loss.mean().item()\n",
    "            print('saving model with loss =', min_loss)\n",
    "            torch.save(model.state_dict(), 'saved_model/decoder2.pth')\n",
    "        loss.mean().backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # Print training statistics\n",
    "        if i % print_interval == 0:\n",
    "            print(f\"Epoch [{epoch}/{num_epochs}], Step [{i}/{len(train_dataloader)}], Loss: {loss.mean().item():.4f}\")\n",
    "\n",
    "    # Update learning rate at the end of each epoch\n",
    "    scheduler.step()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gpu2.0",
   "language": "python",
   "name": "gpu2.0"
  },
  "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}