EE_cifar10_training_2exits.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7f3456be2130>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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.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_2exits_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.49971\n",
      "[2,   100] loss: 0.92502\n",
      "[3,   100] loss: 0.80400\n",
      "[4,   100] loss: 0.72906\n",
      "[5,   100] loss: 0.68492\n",
      "[6,   100] loss: 0.62914\n",
      "[7,   100] loss: 0.60004\n",
      "[8,   100] loss: 0.58075\n",
      "[9,   100] loss: 0.55643\n",
      "[10,   100] loss: 0.53000\n",
      "[11,   100] loss: 0.51498\n",
      "[12,   100] loss: 0.49274\n",
      "[13,   100] loss: 0.49444\n",
      "[14,   100] loss: 0.48575\n",
      "[15,   100] loss: 0.45953\n",
      "[16,   100] loss: 0.44847\n",
      "[17,   100] loss: 0.44738\n",
      "[18,   100] loss: 0.43706\n",
      "[19,   100] loss: 0.42498\n",
      "[20,   100] loss: 0.42027\n",
      "[21,   100] loss: 0.41519\n",
      "[22,   100] loss: 0.40127\n",
      "[23,   100] loss: 0.40035\n",
      "[24,   100] loss: 0.39643\n",
      "[25,   100] loss: 0.38635\n",
      "0.05\n",
      "[26,   100] loss: 0.30852\n",
      "[27,   100] loss: 0.27866\n",
      "[28,   100] loss: 0.26337\n",
      "[29,   100] loss: 0.26368\n",
      "[30,   100] loss: 0.26128\n",
      "[31,   100] loss: 0.25402\n",
      "[32,   100] loss: 0.24306\n",
      "[33,   100] loss: 0.25393\n",
      "[34,   100] loss: 0.24399\n",
      "[35,   100] loss: 0.24395\n",
      "[36,   100] loss: 0.23790\n",
      "[37,   100] loss: 0.23172\n",
      "[38,   100] loss: 0.23258\n",
      "[39,   100] loss: 0.21785\n",
      "[40,   100] loss: 0.21862\n",
      "[41,   100] loss: 0.21936\n",
      "[42,   100] loss: 0.22206\n",
      "[43,   100] loss: 0.21227\n",
      "[44,   100] loss: 0.21705\n",
      "[45,   100] loss: 0.22062\n",
      "[46,   100] loss: 0.21316\n",
      "[47,   100] loss: 0.21096\n",
      "[48,   100] loss: 0.20809\n",
      "[49,   100] loss: 0.20393\n",
      "[50,   100] loss: 0.20179\n",
      "0.025\n",
      "[51,   100] loss: 0.15954\n",
      "[52,   100] loss: 0.13808\n",
      "[53,   100] loss: 0.12368\n",
      "[54,   100] loss: 0.11793\n",
      "[55,   100] loss: 0.11537\n",
      "[56,   100] loss: 0.11084\n",
      "[57,   100] loss: 0.10871\n",
      "[58,   100] loss: 0.11225\n",
      "[59,   100] loss: 0.10897\n",
      "[60,   100] loss: 0.10175\n",
      "[61,   100] loss: 0.10556\n",
      "[62,   100] loss: 0.09956\n",
      "[63,   100] loss: 0.10077\n",
      "[64,   100] loss: 0.09267\n",
      "[65,   100] loss: 0.09848\n",
      "[66,   100] loss: 0.09940\n",
      "[67,   100] loss: 0.10318\n",
      "[68,   100] loss: 0.09734\n",
      "[69,   100] loss: 0.10269\n",
      "[70,   100] loss: 0.09859\n",
      "[71,   100] loss: 0.09163\n",
      "[72,   100] loss: 0.09460\n",
      "[73,   100] loss: 0.09677\n",
      "[74,   100] loss: 0.09478\n",
      "[75,   100] loss: 0.09815\n",
      "0.0125\n",
      "[76,   100] loss: 0.07284\n",
      "[77,   100] loss: 0.05846\n",
      "[78,   100] loss: 0.05414\n",
      "[79,   100] loss: 0.04938\n",
      "[80,   100] loss: 0.04768\n",
      "[81,   100] loss: 0.04810\n",
      "[82,   100] loss: 0.04541\n",
      "[83,   100] loss: 0.04595\n",
      "[84,   100] loss: 0.04856\n",
      "[85,   100] loss: 0.04421\n",
      "[86,   100] loss: 0.04206\n",
      "[87,   100] loss: 0.03968\n",
      "[88,   100] loss: 0.04242\n",
      "[89,   100] loss: 0.03958\n",
      "[90,   100] loss: 0.04231\n",
      "[91,   100] loss: 0.04023\n",
      "[92,   100] loss: 0.03780\n",
      "[93,   100] loss: 0.03753\n",
      "[94,   100] loss: 0.03964\n",
      "[95,   100] loss: 0.04142\n",
      "[96,   100] loss: 0.04193\n",
      "[97,   100] loss: 0.03696\n",
      "[98,   100] loss: 0.03895\n",
      "[99,   100] loss: 0.03887\n",
      "[100,   100] loss: 0.03695\n",
      "0.00625\n",
      "[101,   100] loss: 0.03062\n",
      "[102,   100] loss: 0.02579\n",
      "[103,   100] loss: 0.02566\n",
      "[104,   100] loss: 0.02627\n",
      "[105,   100] loss: 0.02359\n",
      "[106,   100] loss: 0.02312\n",
      "[107,   100] loss: 0.02309\n",
      "[108,   100] loss: 0.02103\n",
      "[109,   100] loss: 0.02072\n",
      "[110,   100] loss: 0.02064\n",
      "[111,   100] loss: 0.02050\n",
      "[112,   100] loss: 0.02027\n",
      "[113,   100] loss: 0.01954\n",
      "[114,   100] loss: 0.02089\n",
      "[115,   100] loss: 0.01852\n",
      "[116,   100] loss: 0.02035\n",
      "[117,   100] loss: 0.01836\n",
      "[118,   100] loss: 0.01859\n",
      "[119,   100] loss: 0.01785\n",
      "[120,   100] loss: 0.01765\n",
      "[121,   100] loss: 0.01748\n",
      "[122,   100] loss: 0.01865\n",
      "[123,   100] loss: 0.01838\n",
      "[124,   100] loss: 0.01718\n",
      "[125,   100] loss: 0.02069\n",
      "0.003125\n",
      "[126,   100] loss: 0.01578\n",
      "[127,   100] loss: 0.01449\n",
      "[128,   100] loss: 0.01405\n",
      "[129,   100] loss: 0.01352\n",
      "[130,   100] loss: 0.01438\n",
      "[131,   100] loss: 0.01244\n",
      "[132,   100] loss: 0.01322\n",
      "[133,   100] loss: 0.01203\n",
      "[134,   100] loss: 0.01299\n",
      "[135,   100] loss: 0.01242\n",
      "[136,   100] loss: 0.01382\n",
      "[137,   100] loss: 0.01093\n",
      "[138,   100] loss: 0.01223\n",
      "[139,   100] loss: 0.01245\n",
      "[140,   100] loss: 0.01291\n",
      "[141,   100] loss: 0.01315\n",
      "[142,   100] loss: 0.01226\n",
      "[143,   100] loss: 0.01195\n",
      "[144,   100] loss: 0.01165\n",
      "[145,   100] loss: 0.01149\n",
      "[146,   100] loss: 0.01159\n",
      "[147,   100] loss: 0.01134\n",
      "[148,   100] loss: 0.01081\n",
      "[149,   100] loss: 0.01208\n",
      "[150,   100] loss: 0.01122\n",
      "Finished Training\n",
      "CPU times: user 36min 22s, sys: 55.7 s, total: 37min 18s\n",
      "Wall time: 37min 41s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "epochs= 150\n",
    "coefficient = [0.3,0.7]\n",
    "loss = []\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, = 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\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_2_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
}