mnist_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self, input_shape=(1, 28, 28)):\n",
    "        super(Net, self).__init__()\n",
    "        self.features = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, 3), nn.ReLU(),\n",
    "            nn.Conv2d(32, 32, 3),nn.ReLU(),\n",
    "            nn.MaxPool2d((2,2)),\n",
    "            nn.Conv2d(32, 64, 3), nn.ReLU(),\n",
    "            nn.Conv2d(64, 64, 3),nn.ReLU(),\n",
    "            nn.MaxPool2d((2,2)),\n",
    "        )\n",
    "        \n",
    "        self.feature_dim = self.get_features_dim(input_shape)\n",
    "        \n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Linear(self.feature_dim, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=0.2),\n",
    "            nn.Linear(128, 10),\n",
    "            nn.LogSoftmax()\n",
    "        )\n",
    "    \n",
    "    def get_features_dim(self, shape):\n",
    "        features = self.features(Variable(torch.ones(1, *shape)))\n",
    "        return int(np.prod(features.size()[1:]))\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.features(x)\n",
    "        x = x.view(-1, self.feature_dim)\n",
    "        x = self.classifier(x)\n",
    "        return F.log_softmax(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Net (\n",
       "  (features): Sequential (\n",
       "    (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (1): ReLU ()\n",
       "    (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (3): ReLU ()\n",
       "    (4): MaxPool2d (size=(2, 2), stride=(2, 2), dilation=(1, 1))\n",
       "    (5): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (6): ReLU ()\n",
       "    (7): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (8): ReLU ()\n",
       "    (9): MaxPool2d (size=(2, 2), stride=(2, 2), dilation=(1, 1))\n",
       "  )\n",
       "  (classifier): Sequential (\n",
       "    (0): Linear (1024 -> 128)\n",
       "    (1): ReLU ()\n",
       "    (2): Dropout (p = 0.2)\n",
       "    (3): Linear (128 -> 10)\n",
       "    (4): LogSoftmax ()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = Net()\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "kwargs = {'num_workers': 1, 'pin_memory': True}\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST('data', train=True, download=True,\n",
    "                   transform=transforms.Compose([\n",
    "                       transforms.ToTensor(),\n",
    "                       transforms.Normalize((0.1307,), (0.3081,))\n",
    "                   ])),\n",
    "    batch_size=64, shuffle=True, **kwargs)\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST('data', train=False, transform=transforms.Compose([\n",
    "                       transforms.ToTensor(),\n",
    "                       transforms.Normalize((0.1307,), (0.3081,))\n",
    "                   ])),\n",
    "    batch_size=64, shuffle=True, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Smooth:\n",
    "    def __init__(self, windowsize=10):\n",
    "        self.window_size = windowsize\n",
    "        self.data = np.zeros((self.window_size), dtype=np.float32)\n",
    "        self.index = 0\n",
    "    \n",
    "    def __iadd__(self, x):\n",
    "        if self.index == 0:\n",
    "            self.data[:] = x\n",
    "        self.data[self.index % self.window_size] = x\n",
    "        self.index += 1\n",
    "        return self\n",
    "    \n",
    "    def __float__(self):\n",
    "        return float(self.data.mean())\n",
    "    \n",
    "    def __str__(self):\n",
    "        return str(float(self))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def test():\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    for data, target in test_loader:\n",
    "        data, target = data.cuda(), target.cuda()\n",
    "        data, target = Variable(data, volatile=True), Variable(target)\n",
    "        output = model(data)\n",
    "        test_loss += F.nll_loss(output, target).data[0]\n",
    "        pred = output.data.max(1)[1]\n",
    "        correct += pred.eq(target.data).cpu().sum()\n",
    "    \n",
    "    test_loss /= len(test_loader)\n",
    "    test_acc = 100. * correct / len(test_loader.dataset)\n",
    "    return test_loss, test_acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from sys import stdout\n",
    "class BarMessage:\n",
    "    def __init__(self, epoch=10, bar_length=25, sep=' - '):\n",
    "        self.epoch = epoch\n",
    "        self.bar_length = bar_length\n",
    "        self.sep = sep\n",
    "        self.start = time.time()\n",
    "        self.last = 0\n",
    "\n",
    "    def message(self, i_epoch, i_batch, epoch, loss, val_loss=None, val_acc=None):\n",
    "        progress = float(i_batch+1)/len(train_loader)\n",
    "        block = int(progress*self.bar_length)\n",
    "\n",
    "        if self.last - self.start < 1000:\n",
    "            self.last = time.time()\n",
    "            self.eta = (1-progress)/progress*(self.last - self.start)\n",
    "        \n",
    "        message = '\\r'\n",
    "        message += 'Epoch %2d/%2d ' % (i_epoch+1, self.epoch)\n",
    "        if progress != 1:\n",
    "            message += '[%s>%s]%s' % ('='*block, '.'*(self.bar_length-block), self.sep)\n",
    "            message += 'ETA: %ds%s'% (self.eta, self.sep)\n",
    "        else:\n",
    "            message += '[%s]%s' % ('='*(self.bar_length+1), self.sep)\n",
    "            message += '%ds%s'%(time.time() - self.start, self.sep)\n",
    "        message += 'loss: %.4f ' % (float(loss))\n",
    "        \n",
    "        if val_loss:\n",
    "            message += 'val_loss: %.4f ' % val_loss\n",
    "        if val_acc:\n",
    "            message += 'val_acc: %.2f ' % val_acc\n",
    "        \n",
    "        message += ' '*5\n",
    "\n",
    "        return message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for a in test_loader:\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch  1/10 [==========================] - 5s - loss: 0.0707 val_loss: 0.0356 val_acc: 98.93      \n",
      "Epoch  2/10 [==========================] - 5s - loss: 0.0434 val_loss: 0.0365 val_acc: 98.94      \n",
      "Epoch  3/10 [==========================] - 4s - loss: 0.0382 val_loss: 0.0286 val_acc: 99.18      \n",
      "Epoch  4/10 [==========================] - 5s - loss: 0.0262 val_loss: 0.0216 val_acc: 99.29      \n",
      "Epoch  5/10 [==========================] - 4s - loss: 0.0203 val_loss: 0.0233 val_acc: 99.35      \n",
      "Epoch  6/10 [==========================] - 5s - loss: 0.0131 val_loss: 0.0262 val_acc: 99.30      \n",
      "Epoch  7/10 [==========================] - 5s - loss: 0.0172 val_loss: 0.0226 val_acc: 99.30      \n",
      "Epoch  8/10 [==========================] - 5s - loss: 0.0135 val_loss: 0.0268 val_acc: 99.24      \n",
      "Epoch  9/10 [==========================] - 5s - loss: 0.0112 val_loss: 0.0242 val_acc: 99.30      \n",
      "Epoch 10/10 [==========================] - 5s - loss: 0.0142 val_loss: 0.0251 val_acc: 99.29      \n"
     ]
    }
   ],
   "source": [
    "epoch = 10\n",
    "for i_epoch in range(epoch):\n",
    "    smooth_loss = Smooth(100)\n",
    "    bar_message = BarMessage(epoch)\n",
    "    model.train()\n",
    "    optimizer = optim.Adam(model.parameters())\n",
    "    for i_batch, (data, target) in enumerate(train_loader):\n",
    "        data, target = Variable(data.cuda()), Variable(target.cuda())\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        smooth_loss += loss.data[0]\n",
    "        \n",
    "        stdout.write(bar_message.message(i_epoch, i_batch, epoch, smooth_loss))\n",
    "        stdout.flush()\n",
    "    \n",
    "    val_loss, val_acc = test()\n",
    "    stdout.write(bar_message.message(i_epoch, i_batch, epoch, \n",
    "                                     smooth_loss, val_loss, val_acc))\n",
    "    print ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}