pytorch/spatial_transformer_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "空间转换网络 (Spatial Transformer Networks) 教程\n",
    "=====================================\n",
    "**原作者**: `Ghassen HAMROUNI <https://github.com/GHamrouni>`_\n",
    "\n",
    ".. figure:: /_static/img/stn/FSeq.png\n",
    "\n",
    "在这篇教程中, 你会学到如何用名为空间转换网络 (spatial transformer networks) \n",
    "的视觉注意力结构来加强你的网络. 你可以从这篇论文上看到更多关于空间转换网络 (spatial \n",
    "transformer networks)的知识: `DeepMind paper <https://arxiv.org/abs/1506.02025>`__\n",
    "\n",
    "空间转换网络 (spatial transformer networks) 是对关注空间变换可区分性的一种推广\n",
    "形式. 短空间转换网络 (STN for short) 允许一个神经网络学习如何在输入图像上表现出空\n",
    "间变换, 以此来增强模型的几何不变性.\n",
    "例如, 它可以裁剪一个感兴趣的区域, 缩放和修正图像的方向. 由于卷积神经网络对旋转、缩放\n",
    "和更普遍仿射变换并不具有不变性, 因此它相对来说是一种有用的结构. \n",
    "\n",
    "STN (空间转换网络) 最好的一点是它能在非常小的改动之后, 被简单地嵌入到任何已存在的卷积神\n",
    "经网络中. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 许可协议: BSD\n",
    "# 作者: Ghassen Hamrouni\n",
    "\n",
    "from __future__ import print_function\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "from torchvision import datasets, transforms\n",
    "from torch.autograd import Variable\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "plt.ion()   # 交互模式"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "读数据\n",
    "----------------\n",
    "\n",
    "在这里我们用经典的 MNIST 数据集做试验. 使用一个被空间转换网络增强的标准卷积神经\n",
    "网络.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
      "Processing...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "use_cuda = torch.cuda.is_available()\n",
    "\n",
    "# 训练集\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST(root='.', train=True, download=True,\n",
    "                   transform=transforms.Compose([\n",
    "                       transforms.ToTensor(),\n",
    "                       transforms.Normalize((0.1307,), (0.3081,))\n",
    "                   ])), batch_size=64, shuffle=True, num_workers=4)\n",
    "# 测试集\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST(root='.', train=False, transform=transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.1307,), (0.3081,))\n",
    "    ])), batch_size=64, shuffle=True, num_workers=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "描述空间转换网络 (spatial transformer networks)\n",
    "--------------------------------------\n",
    "\n",
    "空间转换网络 (spatial transformer networks) 归纳为三个主要的部件 :\n",
    "\n",
    "-  本地网络 (The localization network) 是一个常规CNN, 它可以回归转换参数. \n",
    "   这种空间转换不是简单地从数据集显式学习到的, 而是自动地学习以增强全局准确率.\n",
    "-  网格生成器 (The grid generator) 在输入图像中生成对应于来自输出图像的每个像\n",
    "   素的坐标网格. \n",
    "-  采样器 (The sampler) 将转换的参数应用于输入图像. \n",
    "\n",
    ".. figure:: /_static/img/stn/stn-arch.png\n",
    "\n",
    ".. Note::\n",
    "   我们需要包含 affine_grid 和 grid_sample 模块的 PyTorch 最新版本. \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n",
    "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
    "        self.conv2_drop = nn.Dropout2d()\n",
    "        self.fc1 = nn.Linear(320, 50)\n",
    "        self.fc2 = nn.Linear(50, 10)\n",
    "\n",
    "        # 空间转换本地网络 (Spatial transformer localization-network)\n",
    "        self.localization = nn.Sequential(\n",
    "            nn.Conv2d(1, 8, kernel_size=7),\n",
    "            nn.MaxPool2d(2, stride=2),\n",
    "            nn.ReLU(True),\n",
    "            nn.Conv2d(8, 10, kernel_size=5),\n",
    "            nn.MaxPool2d(2, stride=2),\n",
    "            nn.ReLU(True)\n",
    "        )\n",
    "\n",
    "        # 3 * 2 仿射矩阵 (affine matrix) 的回归器\n",
    "        self.fc_loc = nn.Sequential(\n",
    "            nn.Linear(10 * 3 * 3, 32),\n",
    "            nn.ReLU(True),\n",
    "            nn.Linear(32, 3 * 2)\n",
    "        )\n",
    "\n",
    "        # 用身份转换 (identity transformation) 初始化权重 (weights) / 偏置 (bias)\n",
    "        self.fc_loc[2].weight.data.fill_(0)\n",
    "        self.fc_loc[2].bias.data = torch.FloatTensor([1, 0, 0, 0, 1, 0])\n",
    "\n",
    "    # 空间转换网络的前向函数 (Spatial transformer network forward function)\n",
    "    def stn(self, x):\n",
    "        xs = self.localization(x)\n",
    "        xs = xs.view(-1, 10 * 3 * 3)\n",
    "        theta = self.fc_loc(xs)\n",
    "        theta = theta.view(-1, 2, 3)\n",
    "\n",
    "        grid = F.affine_grid(theta, x.size())\n",
    "        x = F.grid_sample(x, grid)\n",
    "\n",
    "        return x\n",
    "\n",
    "    def forward(self, x):\n",
    "        # 转换输入\n",
    "        x = self.stn(x)\n",
    "\n",
    "        # 执行常规的正向传递\n",
    "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
    "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
    "        x = x.view(-1, 320)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x, dim=1)\n",
    "\n",
    "\n",
    "model = Net()\n",
    "if use_cuda:\n",
    "    model.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "训练模型\n",
    "------------------\n",
    "\n",
    "现在, 让我们用 SGD 算法来训练模型. 这个网络用监督学习的方式学习分类任务. 同时, \n",
    "这个模型以端到端的方式自动地学习空间转换网络 (STN) .\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.SGD(model.parameters(), lr=0.01)\n",
    "\n",
    "\n",
    "def train(epoch):\n",
    "    model.train()\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        if use_cuda:\n",
    "            data, target = data.cuda(), target.cuda()\n",
    "\n",
    "        data, target = Variable(data), Variable(target)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if batch_idx % 500 == 0:\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch, batch_idx * len(data), len(train_loader.dataset),\n",
    "                100. * batch_idx / len(train_loader), loss.data[0]))\n",
    "#\n",
    "# 一个简单的测试程序来测量空间转换网络 (STN) 在 MNIST 上的表现.\n",
    "#\n",
    "\n",
    "\n",
    "def test():\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    for data, target in test_loader:\n",
    "        if use_cuda:\n",
    "            data, target = data.cuda(), target.cuda()\n",
    "        data, target = Variable(data, volatile=True), Variable(target)\n",
    "        output = model(data)\n",
    "\n",
    "        # 累加批loss\n",
    "        test_loss += F.nll_loss(output, target, size_average=False).data[0]\n",
    "        # 得到最大对数几率 (log-probability) 的索引.\n",
    "        pred = output.data.max(1, keepdim=True)[1]\n",
    "        correct += pred.eq(target.data.view_as(pred)).cpu().sum()\n",
    "\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'\n",
    "          .format(test_loss, correct, len(test_loader.dataset),\n",
    "                  100. * correct / len(test_loader.dataset)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可视化空间转换网络 (STN) 的结果\n",
    "---------------------------\n",
    "\n",
    "现在, 我们要检查学到的视觉注意力机制的结果. \n",
    "\n",
    "我们定义一个小的辅助函数, 以在训练过程中可视化转换过程. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "F:\\app\\anaconda\\lib\\site-packages\\ipykernel_launcher.py:19: 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": [
      "Train Epoch: 1 [0/60000 (0%)]\tLoss: 2.288913\n",
      "Train Epoch: 1 [32000/60000 (53%)]\tLoss: 0.547555\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "F:\\app\\anaconda\\lib\\site-packages\\ipykernel_launcher.py:32: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n",
      "F:\\app\\anaconda\\lib\\site-packages\\ipykernel_launcher.py:36: 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": [
      "\n",
      "Test set: Average loss: 0.2396, Accuracy: 9323/10000 (93%)\n",
      "\n",
      "Train Epoch: 2 [0/60000 (0%)]\tLoss: 0.406553\n",
      "Train Epoch: 2 [32000/60000 (53%)]\tLoss: 0.324025\n",
      "\n",
      "Test set: Average loss: 0.1736, Accuracy: 9474/10000 (94%)\n",
      "\n",
      "Train Epoch: 3 [0/60000 (0%)]\tLoss: 0.414806\n",
      "Train Epoch: 3 [32000/60000 (53%)]\tLoss: 0.172294\n",
      "\n",
      "Test set: Average loss: 0.0869, Accuracy: 9729/10000 (97%)\n",
      "\n",
      "Train Epoch: 4 [0/60000 (0%)]\tLoss: 0.150663\n",
      "Train Epoch: 4 [32000/60000 (53%)]\tLoss: 0.293290\n",
      "\n",
      "Test set: Average loss: 0.0762, Accuracy: 9763/10000 (97%)\n",
      "\n",
      "Train Epoch: 5 [0/60000 (0%)]\tLoss: 0.209333\n",
      "Train Epoch: 5 [32000/60000 (53%)]\tLoss: 0.227868\n",
      "\n",
      "Test set: Average loss: 0.0733, Accuracy: 9755/10000 (97%)\n",
      "\n",
      "Train Epoch: 6 [0/60000 (0%)]\tLoss: 0.110285\n",
      "Train Epoch: 6 [32000/60000 (53%)]\tLoss: 0.143455\n",
      "\n",
      "Test set: Average loss: 0.0646, Accuracy: 9807/10000 (98%)\n",
      "\n",
      "Train Epoch: 7 [0/60000 (0%)]\tLoss: 0.417222\n",
      "Train Epoch: 7 [32000/60000 (53%)]\tLoss: 0.163063\n",
      "\n",
      "Test set: Average loss: 0.0514, Accuracy: 9834/10000 (98%)\n",
      "\n",
      "Train Epoch: 8 [0/60000 (0%)]\tLoss: 0.139694\n",
      "Train Epoch: 8 [32000/60000 (53%)]\tLoss: 0.122371\n"
     ]
    }
   ],
   "source": [
    "def convert_image_np(inp):\n",
    "    \"\"\"Convert a Tensor to numpy image.\"\"\"\n",
    "    inp = inp.numpy().transpose((1, 2, 0))\n",
    "    mean = np.array([0.485, 0.456, 0.406])\n",
    "    std = np.array([0.229, 0.224, 0.225])\n",
    "    inp = std * inp + mean\n",
    "    inp = np.clip(inp, 0, 1)\n",
    "    return inp\n",
    "\n",
    "# 我们想要在训练之后可视化空间转换层 (spatial transformers layer) 的输出, 我们\n",
    "# 用 STN 可视化一批输入图像和相对于的转换后的数据. \n",
    "\n",
    "\n",
    "def visualize_stn():\n",
    "    # 得到一批输入数据\n",
    "    data, _ = next(iter(test_loader))\n",
    "    data = Variable(data, volatile=True)\n",
    "\n",
    "    if use_cuda:\n",
    "        data = data.cuda()\n",
    "\n",
    "    input_tensor = data.cpu().data\n",
    "    transformed_input_tensor = model.stn(data).cpu().data\n",
    "\n",
    "    in_grid = convert_image_np(\n",
    "        torchvision.utils.make_grid(input_tensor))\n",
    "\n",
    "    out_grid = convert_image_np(\n",
    "        torchvision.utils.make_grid(transformed_input_tensor))\n",
    "\n",
    "    # 并行地 (side-by-side) 画出结果\n",
    "    f, axarr = plt.subplots(1, 2)\n",
    "    axarr[0].imshow(in_grid)\n",
    "    axarr[0].set_title('Dataset Images')\n",
    "\n",
    "    axarr[1].imshow(out_grid)\n",
    "    axarr[1].set_title('Transformed Images')\n",
    "\n",
    "\n",
    "for epoch in range(1, 20 + 1):\n",
    "    train(epoch)\n",
    "    test()\n",
    "\n",
    "# 在一些输入批次中可视化空间转换网络 (STN) 的转换\n",
    "visualize_stn()\n",
    "\n",
    "plt.ioff()\n",
    "plt.show()"
   ]
  }
 ],
 "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": 1
}