EE_cifar10_training-3exits.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch\n",
    "from utils import *\n",
    "import numpy as np\n",
    "#torch.manual_seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "transform = transforms.Compose([\n",
    "    #transforms.RandomCrop(32, padding=4),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.RandomResizedCrop(32),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),\n",
    "])\n",
    "transform_test = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),\n",
    "])\n",
    "trainset = torchvision.datasets.CIFAR10(root='./data', train=True,\n",
    "                                        download=True, transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=256,\n",
    "                                          shuffle=True, num_workers=8)\n",
    "\n",
    "testset = torchvision.datasets.CIFAR10(root='./data', train=False,\n",
    "                                       download=True, transform=transform_test)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=256,\n",
    "                                         shuffle=False, num_workers=8)\n",
    "\n",
    "classes = ('plane', 'car', 'bird', 'cat',\n",
    "           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import torch.nn as nn\n",
    "\n",
    "model = Resnet18_3exits_Pyramidpool().cuda()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.1,\n",
    "                      momentum=0.6, weight_decay=5e-4)\n",
    "my_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=optimizer, gamma=0.5)\n",
    "#model, optimizer = amp.initialize(model, optimizer, opt_level=\"O2\")\n",
    "#lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer,\n",
    "#                                                              T_0=int(len(trainset)/64*2),\n",
    "     #                                                         T_mult=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,   100] loss: 1.82402\n",
      "[2,   100] loss: 1.40945\n",
      "[3,   100] loss: 1.29109\n",
      "[4,   100] loss: 1.23568\n",
      "[5,   100] loss: 1.17070\n",
      "[6,   100] loss: 1.13242\n",
      "[7,   100] loss: 1.12796\n",
      "[8,   100] loss: 1.08944\n",
      "[9,   100] loss: 1.08144\n",
      "[10,   100] loss: 1.05712\n",
      "[11,   100] loss: 1.04186\n",
      "[12,   100] loss: 1.02549\n",
      "[13,   100] loss: 1.02361\n",
      "[14,   100] loss: 1.01432\n",
      "[15,   100] loss: 1.00912\n",
      "[16,   100] loss: 0.99675\n",
      "[17,   100] loss: 0.98728\n",
      "[18,   100] loss: 0.98989\n",
      "[19,   100] loss: 0.98675\n",
      "[20,   100] loss: 0.96251\n",
      "[21,   100] loss: 0.96108\n",
      "[22,   100] loss: 0.96919\n",
      "[23,   100] loss: 0.96161\n",
      "[24,   100] loss: 0.94671\n",
      "[25,   100] loss: 0.96048\n",
      "0.05\n",
      "[26,   100] loss: 0.90146\n",
      "[27,   100] loss: 0.87557\n",
      "[28,   100] loss: 0.86510\n",
      "[29,   100] loss: 0.86203\n",
      "[30,   100] loss: 0.86229\n",
      "[31,   100] loss: 0.85275\n",
      "[32,   100] loss: 0.85356\n",
      "[33,   100] loss: 0.84624\n",
      "[34,   100] loss: 0.84859\n",
      "[35,   100] loss: 0.84601\n",
      "[36,   100] loss: 0.84981\n",
      "[37,   100] loss: 0.84263\n",
      "[38,   100] loss: 0.84499\n",
      "[39,   100] loss: 0.83951\n",
      "[40,   100] loss: 0.83043\n",
      "[41,   100] loss: 0.83660\n",
      "[42,   100] loss: 0.83229\n",
      "[43,   100] loss: 0.83576\n",
      "[44,   100] loss: 0.83127\n",
      "[45,   100] loss: 0.82841\n",
      "[46,   100] loss: 0.81902\n",
      "[47,   100] loss: 0.82895\n",
      "[48,   100] loss: 0.82334\n",
      "[49,   100] loss: 0.82759\n",
      "[50,   100] loss: 0.82757\n",
      "0.025\n",
      "[51,   100] loss: 0.78638\n",
      "[52,   100] loss: 0.76885\n",
      "[53,   100] loss: 0.76759\n",
      "[54,   100] loss: 0.75521\n",
      "[55,   100] loss: 0.74492\n",
      "[56,   100] loss: 0.75366\n",
      "[57,   100] loss: 0.76110\n",
      "[58,   100] loss: 0.75229\n",
      "[59,   100] loss: 0.75061\n",
      "[60,   100] loss: 0.74682\n",
      "[61,   100] loss: 0.75064\n",
      "[62,   100] loss: 0.74334\n",
      "[63,   100] loss: 0.75278\n",
      "[64,   100] loss: 0.73451\n",
      "[65,   100] loss: 0.74976\n",
      "[66,   100] loss: 0.73792\n",
      "[67,   100] loss: 0.74002\n",
      "[68,   100] loss: 0.74829\n",
      "[69,   100] loss: 0.73542\n",
      "[70,   100] loss: 0.72852\n",
      "[71,   100] loss: 0.74250\n",
      "[72,   100] loss: 0.73086\n",
      "[73,   100] loss: 0.73075\n",
      "[74,   100] loss: 0.73618\n",
      "[75,   100] loss: 0.73803\n",
      "Finished Training\n",
      "CPU times: user 23min 28s, sys: 24.9 s, total: 23min 53s\n",
      "Wall time: 24min 5s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "epochs= 75\n",
    "loss = []\n",
    "coefficient = [0.2,0.3,0.5]\n",
    "for epoch in range(epochs):\n",
    "    if epoch%25==0 and epoch!=0:\n",
    "        my_lr_scheduler.step()\n",
    "        print(optimizer.param_groups[0]['lr'])\n",
    "        \n",
    "        # loop over the dataset multiple times\n",
    "\n",
    "    running_loss = 0.0\n",
    "    for i, data in enumerate(trainloader, 0):\n",
    "        # get the inputs; data is a list of [inputs, labels]\n",
    "        inputs, labels = data\n",
    "        inputs = inputs.cuda()\n",
    "        labels = labels.cuda()\n",
    "\n",
    "        # zero the parameter gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # forward + backward + optimize\n",
    "        exit1,exit2,exit3 = model(inputs)\n",
    "        loss1 = criterion(exit1, labels)\n",
    "        loss2 = criterion(exit2, labels)\n",
    "        loss3 = criterion(exit3, labels)\n",
    "        total_loss = coefficient[0]*loss1+coefficient[1]*loss2+coefficient[2]*loss3\n",
    "       # with amp.scale_loss(loss, optimizer) as scaled_loss:\n",
    "        #    scaled_loss.backward()\n",
    "        total_loss.backward()\n",
    "        optimizer.step()\n",
    "        #lr_scheduler.step()\n",
    "        # print statistics\n",
    "        running_loss += total_loss.item()\n",
    "        \n",
    "        \n",
    "        if i % 100 == 99:    # print every 2000 mini-batches\n",
    "            print('[%d, %5d] loss: %.5f' %\n",
    "                  (epoch + 1, i + 1, running_loss /99))\n",
    "            loss.append(running_loss)\n",
    "            running_loss = 0.0\n",
    "\n",
    "    torch.save(model,\"Pyramid_3_exits.pth\")\n",
    "print('Finished Training')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}