train_VGG16.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from network_definitions import VGG16  # or whichever model you want to train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# =================================================== Prepare the dataset ===============================================================================\n",
    "\n",
    "mean_cifar10 = [0.485, 0.456, 0.406]  # Mean and Std value hase been taken from a github implmentation online.\n",
    "std_cifar10 = [0.229, 0.224, 0.225]\n",
    "batch_size = 100\n",
    "transform_train = transforms.Compose([\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean_cifar10,std_cifar10),\n",
    "])\n",
    "\n",
    "transform_test = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean_cifar10,std_cifar10),\n",
    "])\n",
    "\n",
    "\n",
    "trainset = torchvision.datasets.CIFAR10(root='../FlexibleCNNs/data', train=True, download= True, transform=transform_train)\n",
    "\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)\n",
    "\n",
    "testset = torchvision.datasets.CIFAR10(root='../FlexibleCNNs/data', train=False, download=True, transform=transform_test)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)\n",
    "\n",
    "classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') # 10 Classes of the cifar-10\n",
    "\n",
    "# ========================================== Visualising the dataset ==========================================================================\n",
    "std= torch.FloatTensor(std_cifar10)\n",
    "mean = torch.FloatTensor(mean_cifar10)\n",
    "mean = mean[:,None,None]\n",
    "std = std[:,None,None]\n",
    "def imshow(img):\n",
    "    print(img.size())\n",
    "    img = img*std + mean     # unnormalize\n",
    "    \n",
    "    npimg = img.numpy()\n",
    "    \n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "\n",
    "\n",
    "# get some random training images\n",
    "dataiter = iter(trainloader)\n",
    "images, labels = next(dataiter)\n",
    "\n",
    "# show images\n",
    "imshow(torchvision.utils.make_grid(images[:4]))\n",
    "# print labels\n",
    "print(' '.join('%5s' % classes[labels[j]] for j in range(4)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# =============================================================== Model initialisation, Loss function and Optimizer =====================================\n",
    "model = VGG16()\n",
    "if torch.cuda.is_available():\n",
    "    model.cuda()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(model.parameters(),lr = 0.001,momentum = 0.9,weight_decay = 0.006)\n",
    "schedule = torch.optim.lr_scheduler.StepLR(optimizer,step_size=20,gamma = 0.7)\n",
    "\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f'{total_params:,} total parameters.')\n",
    "total_trainable_params = sum(\n",
    "    p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print(f'{total_trainable_params:,} trainable parameters.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# ======================== Function to get the test accuracy ===============================================================================\n",
    "def test():\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    model.train(False)\n",
    "    with torch.no_grad():\n",
    "        for i,(images,labels)in enumerate(testloader):\n",
    "            if torch.cuda.is_available():\n",
    "                images = images.cuda()\n",
    "                labels = labels.cuda()\n",
    "            outputs = model(Variable(images))\n",
    "            labels = Variable(labels)\n",
    "            _,predicted = outputs.max(1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted.eq(labels)).sum().item()\n",
    "        print('Test accuracy: %d %%' % (\n",
    "            100 * correct / total))\n",
    "    return 100*(correct/total)\n",
    "\n",
    "#======================================================= Training =========================================================================\n",
    "num_epochs = 152  # Train for 150 epochs\n",
    "start_epoch = 0\n",
    "\n",
    "total_step = len(trainloader)\n",
    "train_loss = []  # Store the train_loss per epoch\n",
    "test_accuracy = [] # Store the test_accuracy per epoch\n",
    "for epoch in range(start_epoch,num_epochs+1):\n",
    "    model.train(True)\n",
    "    epoch_loss  = 0\n",
    "    i_count = 0\n",
    "    acc_total = 0\n",
    "    for i,(images,labels) in enumerate(trainloader):\n",
    "        if torch.cuda.is_available():\n",
    "            images = images.cuda()\n",
    "            labels = labels.cuda()\n",
    "        labels = Variable(labels)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(Variable(images))\n",
    "        loss = criterion(outputs,labels)\n",
    "        epoch_loss += loss.item()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        _,predicted = outputs.max(1)\n",
    "        denom = labels.size(0)\n",
    "        correct = predicted.eq(labels).sum().item()\n",
    "        acc = 100*(correct/denom)\n",
    "        acc_total += acc\n",
    "        i_count = i_count + 1\n",
    "    \n",
    "    schedule.step()\n",
    "    train_loss.append(epoch_loss)\n",
    "    print(\"Epoch: \",epoch,\" \",\"Loss: \",epoch_loss,\" \",\"Train Accuracy :\",acc_total/i_count) # Print train accuracy per epoch\n",
    "    print('\\n')\n",
    "    test_acc = test()      # Print the test accuracy per epoch\n",
    "    test_accuracy.append(test_acc)\n",
    "  \n",
    "    state = {\n",
    "        'model': model.state_dict(),\n",
    "        'acc' : test_acc,\n",
    "        'optim':optimizer.state_dict(),\n",
    "        'epoch' : epoch\n",
    "    }\n",
    "#     path = f'./models/{model.name}' + 'model_' + str(int(epoch)) +'_' + str(int(test_acc))+'.pth'\n",
    "#     torch.save(state,path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "with open(f\"{model.name}_accuracy.pkl\", \"wb\") as output_file:\n",
    "        pickle.dump(test_accuracy, output_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}