Pytorch_CNN_CIFAR.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data_utils\n",
    "import torch.nn.init as init\n",
    "from torch.autograd import Variable\n",
    "from dataloader import load_cifar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparams\n",
    "EPOCHS  =  3\n",
    "BATCHSIZE  =  64\n",
    "LR  =  0.01\n",
    "MOMENTUM  =  0.9\n",
    "N_CLASSES  =  10\n",
    "GPU  =  True\n",
    "BATCH_SIZE  =  32\n",
    "RESNET_FEATURES  =  2048\n",
    "BATCHES_GPU  =  40\n",
    "\n",
    "torch.backends.cudnn.benchmark=True\n",
    "torch.cuda.current_device()\n",
    "torch.cuda._initialized = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train, x_test, y_train, y_test = load_cifar(channel_first=True)\n",
    "# Torch-specific\n",
    "y_train = y_train.astype(np.int32)\n",
    "y_test = y_test.astype(np.int64)\n",
    "print(x_train.shape, x_test.shape, y_train.shape, y_test.shape)\n",
    "print(x_train.dtype, x_test.dtype, y_train.dtype, y_test.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SymbolModule(nn.Module):\n",
    "    def __init__(self, n_classes=N_CLASSES):\n",
    "        super(SymbolModule, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 50, kernel_size=3, padding=1)\n",
    "        self.conv2 = nn.Conv2d(50, 50, kernel_size=3, padding=1)\n",
    "        self.conv3 = nn.Conv2d(50, 100, kernel_size=3, padding=1)\n",
    "        self.conv4 = nn.Conv2d(100, 100, kernel_size=3, padding=1)\n",
    "        # feature map size is 8*8 by pooling\n",
    "        self.fc1 = nn.Linear(100*8*8, 512)\n",
    "        self.fc2 = nn.Linear(512, n_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # PyTorch requires a flag for training in dropout\n",
    "        x = self.conv2(F.relu(self.conv1(x)))\n",
    "        x = F.relu(F.max_pool2d(x, kernel_size=2, stride=2))\n",
    "        x = F.dropout(x, 0.25, training=self.training)\n",
    "\n",
    "        x = self.conv4(F.relu(self.conv3(x)))\n",
    "        x = F.relu(F.max_pool2d(x, kernel_size=2, stride=2))\n",
    "        x = F.dropout(x, 0.25, training=self.training)\n",
    "\n",
    "        x = x.view(-1, 100*8*8)   # reshape Variable\n",
    "        x = F.dropout(F.relu(self.fc1(x)), 0.5, training=self.training)\n",
    "        return self.fc2(x)\n",
    "    \n",
    "def init_model(m, lr=LR, momentum=MOMENTUM):\n",
    "    # This criterion combines nn.LogSoftmax() and nn.NLLLoss() in one single class\n",
    "    opt = optim.SGD(m.parameters(), lr, momentum)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    return opt, criterion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Init model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sym = SymbolModule()\n",
    "sym.cuda() # CUDA!\n",
    "optimizer, criterion = init_model(sym)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Main training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def shuffle_data(X, y):\n",
    "    s = np.arange(len(X))\n",
    "    np.random.shuffle(s)\n",
    "    X = X[s]\n",
    "    y = y[s]\n",
    "    return X, y\n",
    "\n",
    "def yield_mb(X, y, batchsize=64, shuffle=False):\n",
    "    if shuffle:\n",
    "        X, y = shuffle_data(X, y)\n",
    "    # Only complete batches are submitted\n",
    "    for i in range(len(X) // batchsize):\n",
    "        yield X[i * batchsize:(i + 1) * batchsize], y[i * batchsize:(i + 1) * batchsize]\n",
    "\n",
    "sym.train() # Sets training = True  \n",
    "for j in range(EPOCHS):\n",
    "    for data, target in yield_mb(x_train, y_train, BATCHSIZE, shuffle=True):\n",
    "        # Get samples\n",
    "        data = Variable(torch.FloatTensor(data).cuda())\n",
    "        target = Variable(torch.LongTensor(target).cuda())\n",
    "        # Init\n",
    "        optimizer.zero_grad()\n",
    "        # Forwards\n",
    "        output = sym(data)\n",
    "        # Loss\n",
    "        loss = criterion(output, target)\n",
    "        # Back-prop\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    # Log\n",
    "    print(\"EPOCHS:{},LOSS:{}\".format(j,loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Main evaluation loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sym.eval() # Sets training = False\n",
    "n_samples = (y_test.shape[0]//BATCHSIZE)*BATCHSIZE\n",
    "y_guess = np.zeros(n_samples, dtype=np.int)\n",
    "y_truth = y_test[:n_samples]\n",
    "c = 0\n",
    "for data, target in yield_mb(x_test, y_test, BATCHSIZE):\n",
    "    # Get samples\n",
    "    data = Variable(torch.FloatTensor(data).cuda())\n",
    "    # Forwards\n",
    "    output = sym(data)\n",
    "    pred = output.data.max(1)[1].cpu().numpy().squeeze()\n",
    "    # Collect results\n",
    "    y_guess[c*BATCHSIZE:(c+1)*BATCHSIZE] = pred\n",
    "    c += 1\n",
    "\n",
    "print(\"Accuracy: \", 1.*sum(y_guess == y_truth)/len(y_guess))"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}