MobileNet_CIFAR.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data\n",
    "from torch.autograd import Variable\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "from tensorboardX import SummaryWriter\n",
    "from tqdm import tqdm_notebook\n",
    "\n",
    "import time\n",
    "\n",
    "from nets.MobileNet_CIFAR import MobileNet_CIFAR\n",
    "from utils import build_dataset, flops_to_string, params_to_string, accuracy, file_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Solver(object):\n",
    "    def __init__(self, model, dataset='MNIST', n_epochs=500, lr=0.001, batch_size = 32):\n",
    "        self.n_epochs = n_epochs\n",
    "        \n",
    "        self.train_loader, self.test_loader = build_dataset(dataset, './data', batch_size = batch_size)\n",
    "        \n",
    "        self.writer = SummaryWriter(f'runs/MobileNet_CIFAR-{int(time.time())}')\n",
    "\n",
    "        self.image_dim = {'MNIST': 28*28, 'CIFAR10': 3*32*32}[dataset]\n",
    "                \n",
    "        self.net = model().cuda()\n",
    "        \n",
    "        self.loss_fn = nn.CrossEntropyLoss()\n",
    "        self.optimizer = optim.Adam(self.net.parameters(), lr=lr)\n",
    "                \n",
    "    def train(self):\n",
    "        self.net.train()\n",
    "        self.best_loss = 999\n",
    "        for epoch_i in tqdm_notebook(range(self.n_epochs)):\n",
    "            epoch_loss = 0\n",
    "            for batch_idx, (images, labels) in enumerate(self.train_loader):\n",
    "                images = Variable(images).cuda()\n",
    "                labels = Variable(labels).cuda()\n",
    "                \n",
    "                logits = self.net(images)\n",
    "                \n",
    "                loss = self.loss_fn(logits, labels)\n",
    "                \n",
    "                total_loss = loss\n",
    "\n",
    "                self.optimizer.zero_grad()\n",
    "                total_loss.backward()\n",
    "\n",
    "                self.optimizer.step()\n",
    "                \n",
    "                epoch_loss += float(loss.data)    \n",
    "            epoch_loss /= len(self.train_loader.dataset)\n",
    "            # save best model\n",
    "            if self.best_loss > epoch_loss:\n",
    "                self.best_loss = epoch_loss\n",
    "                torch.save(self.net.state_dict(), f'models/MobileNet_CIFAR_best_loss.pth')\n",
    "                \n",
    "            print(f'Epoch {epoch_i} | loss: {epoch_loss:.8f}')\n",
    "            \n",
    "            niter = epoch_i*len(self.train_loader)+batch_idx\n",
    "            self.writer.add_scalar('Train/Loss', loss.data[0], niter)\n",
    "            self.evaluate(epoch_i)\n",
    "            \n",
    "        # load and test\n",
    "        self.net.load_state_dict(torch.load(f'models/MobileNet_CIFAR_best_loss.pth'))\n",
    "        print (f'Best Loss {self.best_loss} | Its Acc:')\n",
    "        self.evaluate(epoch_i+1)\n",
    "        print (f'Best Top-1 acc over all {self.best_acc1}')\n",
    "        print (f'Best Top-5 acc over all {self.best_acc5}')\n",
    "        \n",
    "        # Print info\n",
    "        pytorch_total_params = sum(p.numel() for p in self.net.parameters())\n",
    "        print (f'Model MobileNet_LowRank | Params = ', params_to_string(pytorch_total_params))\n",
    "        print('Model size: ', file_size(f'models/MobileNet_CIFAR_best_loss.pth'))        \n",
    "        \n",
    "    def evaluate(self, epoch_i = None):\n",
    "        total = 0\n",
    "        correct = 0\n",
    "        self.net.eval()\n",
    "        \n",
    "        self.best_acc1 = 0\n",
    "        self.best_acc5 = 0\n",
    "        acc1_list, acc5_list = [], []\n",
    "        for batch_idx, (images, labels) in enumerate(self.test_loader):\n",
    "            images = Variable(images).cuda()\n",
    "\n",
    "            logits = self.net(images)\n",
    "            _, predicted = torch.max(logits.data, 1)\n",
    "\n",
    "            acc1, acc5 = accuracy(logits.cpu(), labels.cpu(), topk=(1, 5))\n",
    "            \n",
    "            acc1_list.append(acc1)\n",
    "            acc5_list.append(acc5)\n",
    "        \n",
    "        acc1 = float(sum(acc1_list)/len(self.test_loader))\n",
    "        acc5 = float(sum(acc5_list)/len(self.test_loader))\n",
    "        \n",
    "        if self.best_acc1 < acc1:\n",
    "            self.best_acc1 = acc1\n",
    "            torch.save(self.net.state_dict(), f'models/MobileNet_CIFAR_best_acc.pth')\n",
    "        \n",
    "        if self.best_acc5 < acc5:\n",
    "            self.best_acc5 = acc5        \n",
    "        \n",
    "            total += labels.size(0)\n",
    "            correct += (predicted.cpu() == labels).sum()\n",
    "        if epoch_i is not None:\n",
    "            self.writer.add_scalar('Test/Acc@1', acc1, epoch_i)\n",
    "            self.writer.add_scalar('Test/Acc@5', acc5, epoch_i)\n",
    "        print(f'Top-1 Accuracy: {acc1}')\n",
    "        print(f'Top-5 Accuracy: {acc5}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b67761a91d134b7896af4b26f3c05896",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=1), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0 | loss: 0.01331364\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/denis/anaconda3/lib/python3.6/site-packages/ipykernel_launcher.py:46: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Top-1 Accuracy: 42.70174026489258\n",
      "Top-5 Accuracy: 92.11827850341797\n",
      "\n",
      "Best Loss 0.013313643674850464 | Its Acc:\n",
      "Top-1 Accuracy: 42.70174026489258\n",
      "Top-5 Accuracy: 92.11827850341797\n",
      "Model MobileNet_LowRank | Params =  121.61k\n",
      "Model size:  512.3 KB\n",
      "Best Top-1 acc over all 42.70174026489258\n",
      "Best Top-5 acc over all 92.11827850341797\n"
     ]
    }
   ],
   "source": [
    "standard_solver = Solver(MobileNet_CIFAR, 'CIFAR10', n_epochs=500, lr=0.001, batch_size = 128)\n",
    "\n",
    "standard_solver.train()"
   ]
  }
 ],
 "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}