end_to_end_image_compression.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "end_to_end_image_compression.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/kunalrdeshmukh/End-to-end-compression/blob/master/end_to_end_image_compression.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "metadata": {
        "id": "SA29G6TcgrZz",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Image Compression framework :\n",
        "\n",
        "[End to end image compression framework ](https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=7999241)"
      ]
    },
    {
      "metadata": {
        "id": "D98I4TsNVy2W",
        "colab_type": "code",
        "outputId": "9d6e5967-9dbf-472c-a14a-5a5595716a8a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 326
        }
      },
      "cell_type": "code",
      "source": [
        "# http://pytorch.org/\n",
        "from os.path import exists\n",
        "from wheel.pep425tags import get_abbr_impl, get_impl_ver, get_abi_tag\n",
        "platform = '{}{}-{}'.format(get_abbr_impl(), get_impl_ver(), get_abi_tag())\n",
        "cuda_output = !ldconfig -p|grep cudart.so|sed -e 's/.*\\.\\([0-9]*\\)\\.\\([0-9]*\\)$/cu\\1\\2/'\n",
        "accelerator = cuda_output[0] if exists('/dev/nvidia0') else 'cpu'\n",
        "\n",
        "!pip install -q http://download.pytorch.org/whl/{accelerator}/torch-0.4.1-{platform}-linux_x86_64.whl torchvision\n",
        "import torch\n",
        "\n",
        "import os\n",
        "\n",
        "import torch\n",
        "import torchvision\n",
        "from torch import nn , optim\n",
        "from torchvision import datasets, transforms\n",
        "from torch.autograd import Variable\n",
        "\n",
        "!pip install Pillow==4.0.0\n",
        "!pip install PIL\n",
        "!pip install image\n",
        "from PIL import Image"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting Pillow==4.0.0\n",
            "  Using cached https://files.pythonhosted.org/packages/37/e8/b3fbf87b0188d22246678f8cd61e23e31caa1769ebc06f1664e2e5fe8a17/Pillow-4.0.0-cp36-cp36m-manylinux1_x86_64.whl\n",
            "Requirement already satisfied: olefile in /usr/local/lib/python3.6/dist-packages (from Pillow==4.0.0) (0.46)\n",
            "\u001b[31mtorchvision 0.2.1 has requirement pillow>=4.1.1, but you'll have pillow 4.0.0 which is incompatible.\u001b[0m\n",
            "Installing collected packages: Pillow\n",
            "  Found existing installation: Pillow 5.3.0\n",
            "    Uninstalling Pillow-5.3.0:\n",
            "      Successfully uninstalled Pillow-5.3.0\n",
            "Successfully installed Pillow-4.0.0\n",
            "Collecting PIL\n",
            "\u001b[31m  Could not find a version that satisfies the requirement PIL (from versions: )\u001b[0m\n",
            "\u001b[31mNo matching distribution found for PIL\u001b[0m\n",
            "Requirement already satisfied: image in /usr/local/lib/python3.6/dist-packages (1.5.27)\n",
            "Requirement already satisfied: pillow in /usr/local/lib/python3.6/dist-packages (from image) (4.0.0)\n",
            "Requirement already satisfied: django in /usr/local/lib/python3.6/dist-packages (from image) (2.1.3)\n",
            "Requirement already satisfied: olefile in /usr/local/lib/python3.6/dist-packages (from pillow->image) (0.46)\n",
            "Requirement already satisfied: pytz in /usr/local/lib/python3.6/dist-packages (from django->image) (2018.7)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "L9KZr6xwIH9V",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import math\n",
        "irange = range\n",
        "\n",
        "\n",
        "def make_grid(tensor, nrow=8, padding=2,\n",
        "              normalize=False, range=None, scale_each=False, pad_value=0):\n",
        "    \"\"\"Make a grid of images.\n",
        "    Args:\n",
        "        tensor (Tensor or list): 4D mini-batch Tensor of shape (B x C x H x W)\n",
        "            or a list of images all of the same size.\n",
        "        nrow (int, optional): Number of images displayed in each row of the grid.\n",
        "            The Final grid size is (B / nrow, nrow). Default is 8.\n",
        "        padding (int, optional): amount of padding. Default is 2.\n",
        "        normalize (bool, optional): If True, shift the image to the range (0, 1),\n",
        "            by subtracting the minimum and dividing by the maximum pixel value.\n",
        "        range (tuple, optional): tuple (min, max) where min and max are numbers,\n",
        "            then these numbers are used to normalize the image. By default, min and max\n",
        "            are computed from the tensor.\n",
        "        scale_each (bool, optional): If True, scale each image in the batch of\n",
        "            images separately rather than the (min, max) over all images.\n",
        "        pad_value (float, optional): Value for the padded pixels.\n",
        "    Example:\n",
        "        See this notebook `here <https://gist.github.com/anonymous/bf16430f7750c023141c562f3e9f2a91>`_\n",
        "    \"\"\"\n",
        "    if not (torch.is_tensor(tensor) or\n",
        "            (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):\n",
        "        raise TypeError('tensor or list of tensors expected, got {}'.format(type(tensor)))\n",
        "\n",
        "    # if list of tensors, convert to a 4D mini-batch Tensor\n",
        "    if isinstance(tensor, list):\n",
        "        tensor = torch.stack(tensor, dim=0)\n",
        "\n",
        "    if tensor.dim() == 2:  # single image H x W\n",
        "        tensor = tensor.view(1, tensor.size(0), tensor.size(1))\n",
        "    if tensor.dim() == 3:  # single image\n",
        "        if tensor.size(0) == 1:  # if single-channel, convert to 3-channel\n",
        "            tensor = torch.cat((tensor, tensor, tensor), 0)\n",
        "        tensor = tensor.view(1, tensor.size(0), tensor.size(1), tensor.size(2))\n",
        "\n",
        "    if tensor.dim() == 4 and tensor.size(1) == 1:  # single-channel images\n",
        "        tensor = torch.cat((tensor, tensor, tensor), 1)\n",
        "\n",
        "    if normalize is True:\n",
        "        tensor = tensor.clone()  # avoid modifying tensor in-place\n",
        "        if range is not None:\n",
        "            assert isinstance(range, tuple), \\\n",
        "                \"range has to be a tuple (min, max) if specified. min and max are numbers\"\n",
        "\n",
        "        def norm_ip(img, min, max):\n",
        "            img.clamp_(min=min, max=max)\n",
        "            img.add_(-min).div_(max - min + 1e-5)\n",
        "\n",
        "        def norm_range(t, range):\n",
        "            if range is not None:\n",
        "                norm_ip(t, range[0], range[1])\n",
        "            else:\n",
        "                norm_ip(t, float(t.min()), float(t.max()))\n",
        "\n",
        "        if scale_each is True:\n",
        "            for t in tensor:  # loop over mini-batch dimension\n",
        "                norm_range(t, range)\n",
        "        else:\n",
        "            norm_range(tensor, range)\n",
        "\n",
        "    if tensor.size(0) == 1:\n",
        "        return tensor.squeeze()\n",
        "\n",
        "    # make the mini-batch of images into a grid\n",
        "    nmaps = tensor.size(0)\n",
        "    xmaps = min(nrow, nmaps)\n",
        "    ymaps = int(math.ceil(float(nmaps) / xmaps))\n",
        "    height, width = int(tensor.size(2) + padding), int(tensor.size(3) + padding)\n",
        "    grid = tensor.new(3, height * ymaps + padding, width * xmaps + padding).fill_(pad_value)\n",
        "    k = 0\n",
        "    for y in irange(ymaps):\n",
        "        for x in irange(xmaps):\n",
        "            if k >= nmaps:\n",
        "                break\n",
        "            grid.narrow(1, y * height + padding, height - padding)\\\n",
        "                .narrow(2, x * width + padding, width - padding)\\\n",
        "                .copy_(tensor[k])\n",
        "            k = k + 1\n",
        "    return grid\n",
        "\n",
        "\n",
        "def save_image(tensor, filename, nrow=8, padding=2,\n",
        "               normalize=False, range=None, scale_each=False, pad_value=0):\n",
        "    \"\"\"Save a given Tensor into an image file.\n",
        "    Args:\n",
        "        tensor (Tensor or list): Image to be saved. If given a mini-batch tensor,\n",
        "            saves the tensor as a grid of images by calling ``make_grid``.\n",
        "        **kwargs: Other arguments are documented in ``make_grid``.\n",
        "    \"\"\"\n",
        "    from PIL import Image\n",
        "    grid = make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,\n",
        "                     normalize=normalize, range=range, scale_each=scale_each)\n",
        "    ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()\n",
        "    im = Image.fromarray(ndarr)\n",
        "    im.save(filename)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "NG0aiMZcvvcg",
        "colab_type": "code",
        "outputId": "b493dca8-4432-472d-cee4-781f32bf3665",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "cell_type": "code",
      "source": [
        "img_transform = transforms.Compose([\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize((0.5, 0.5, 0.5), (0.24703223,  0.24348513 , 0.26158784))\n",
        "])\n",
        "\n",
        "trainset = datasets.CIFAR10(root='./data', train=True,download=True, transform=img_transform)\n",
        "\n",
        "train_loader = torch.utils.data.DataLoader(trainset, batch_size=16,\n",
        "                                          shuffle=True, num_workers=2)\n",
        "\n",
        "testset = datasets.CIFAR10(root='./data', train=False,download=True, transform=img_transform)\n",
        "\n",
        "test_loader = torch.utils.data.DataLoader(testset, batch_size=16,shuffle=False, num_workers=2)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Files already downloaded and verified\n",
            "Files already downloaded and verified\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "Ggg832Bqwbq7",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# code block to text image loader\n",
        "# print(trainloader)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "KwY8W455pGhk",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Parameters"
      ]
    },
    {
      "metadata": {
        "id": "G3n2__mWHBQI",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "CHANNELS = 3\n",
        "HEIGHT = 32\n",
        "WIDTH = 32\n",
        "EPOCHS = 20\n",
        "LOG_INTERVAL = 500"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "NOVR5jNOU5A7",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Define the network"
      ]
    },
    {
      "metadata": {
        "id": "U5UBMyfgwO5z",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class Interpolate(nn.Module):\n",
        "    def __init__(self, size, mode):\n",
        "        super(Interpolate, self).__init__()\n",
        "        self.interp = nn.functional.interpolate\n",
        "        self.size = size\n",
        "        self.mode = mode\n",
        "        \n",
        "    def forward(self, x):\n",
        "        x = self.interp(x, size=self.size, mode=self.mode, align_corners=False)\n",
        "        return x"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "REJrpjnYFeMm",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "class End_to_end(nn.Module):\n",
        "  def __init__(self):\n",
        "    super(End_to_end, self).__init__()\n",
        "    \n",
        "    # Encoder\n",
        "    # TODO : try with padding = 0\n",
        "    self.conv1 = nn.Conv2d(CHANNELS, 64, kernel_size=3, stride=1, padding=1)\n",
        "    self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=0)\n",
        "    self.bn1 = nn.BatchNorm2d(64, affine=False)\n",
        "    self.conv3 = nn.Conv2d(64, CHANNELS, kernel_size=3, stride=1, padding=1)\n",
        "    \n",
        "    # Decoder\n",
        "    #TODO : try ConvTranspose2d\n",
        "    self.interpolate = Interpolate(size=HEIGHT, mode='bilinear')\n",
        "    self.deconv1 = nn.Conv2d(CHANNELS, 64, 3, stride=1, padding=1)\n",
        "    self.deconv2 = nn.Conv2d(64, 64, 3, stride=1, padding=1)\n",
        "    self.bn2 = nn.BatchNorm2d(64, affine=False)\n",
        "    \n",
        "    self.deconv_n = nn.Conv2d(64, 64, 3, stride=1, padding=1)\n",
        "    self.bn_n = nn.BatchNorm2d(64, affine=False)\n",
        "\n",
        "    \n",
        "    self.deconv3 = nn.ConvTranspose2d(64, CHANNELS, 3, stride=1, padding=1)\n",
        "    \n",
        "    self.relu = nn.ReLU()\n",
        "  \n",
        "  def encode(self, x):\n",
        "    out = self.relu(self.conv1(x))\n",
        "    out = self.relu(self.conv2(out))\n",
        "    out = self.bn1(out)\n",
        "    return self.conv3(out)\n",
        "    \n",
        "  \n",
        "  def reparameterize(self, mu, logvar):\n",
        "    pass\n",
        "  \n",
        "  def decode(self, z):\n",
        "    upscaled_image = self.interpolate(z)\n",
        "    out = self.relu(self.deconv1(upscaled_image))\n",
        "    out = self.relu(self.deconv2(out))\n",
        "    out = self.bn2(out)\n",
        "    for _ in range(10):\n",
        "      out = self.relu(self.deconv_n(out))\n",
        "      out = self.bn_n(out)\n",
        "    out = self.deconv3(out)\n",
        "    final = upscaled_image + out\n",
        "    return final,out,upscaled_image\n",
        "\n",
        "    \n",
        "  def forward(self, x):\n",
        "    com_img = self.encode(x)\n",
        "    final,out,upscaled_image = self.decode(com_img)\n",
        "    return final, out, upscaled_image, com_img, x\n",
        "\n",
        "  \n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "XG45FHw7MFzW",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Optimizer and loss function "
      ]
    },
    {
      "metadata": {
        "id": "RrnHrzxRIP0q",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "CUDA = torch.cuda.is_available()\n",
        "if CUDA:\n",
        "  model = End_to_end().cuda()\n",
        "else :\n",
        "  model = end_to_end()\n",
        "  \n",
        "optimizer = optim.Adam(model.parameters(), lr=1e-3)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "VfZ9DoJogyOA",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def loss_function(final_img,residual_img,upscaled_img,com_img,orig_img):\n",
        "\n",
        "  com_loss = nn.MSELoss(size_average=False)(orig_img, final_img)\n",
        "  rec_loss = nn.MSELoss(size_average=False)(residual_img,orig_img-upscaled_img)\n",
        "  \n",
        "  return com_loss + rec_loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "u9QRvcpIMPGU",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Train , Test functions"
      ]
    },
    {
      "metadata": {
        "id": "dWsyut4Kbruf",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def train(epoch):\n",
        "    model.train()\n",
        "    train_loss = 0\n",
        "    for batch_idx, (data, _) in enumerate(train_loader):\n",
        "        data = Variable(data)\n",
        "        optimizer.zero_grad()\n",
        "        final, residual_img, upscaled_image, com_img, orig_im = model(data.cuda())\n",
        "        loss = loss_function(final, residual_img, upscaled_image, com_img, orig_im)\n",
        "        loss.backward()\n",
        "        train_loss += loss.data[0]\n",
        "        optimizer.step()\n",
        "        if batch_idx % LOG_INTERVAL == 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),\n",
        "                loss.data[0] / len(data)))\n",
        "\n",
        "    print('====> Epoch: {} Average loss: {:.4f}'.format(\n",
        "          epoch, train_loss / len(train_loader.dataset)))\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Mpg44nfscmcG",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "def test(epoch):\n",
        "  \n",
        "  model.eval()\n",
        "  test_loss = 0\n",
        "  for i, (data, _) in enumerate(test_loader):\n",
        "        data = Variable(data, volatile=True)\n",
        "        final, residual_img, upscaled_image, com_img, orig_im = model(data.cuda())\n",
        "        test_loss += loss_function(final, residual_img, upscaled_image, com_img, orig_im).data[0]\n",
        "        if epoch == EPOCHS and i == 0:\n",
        "#             save_image(final.data[0],'reconstruction_final',nrow=8)\n",
        "#             save_image(com_img.data[0],'com_img',nrow=8)\n",
        "            n = min(data.size(0), 6)\n",
        "            print(\"saving the image \"+str(n))\n",
        "            comparison = torch.cat([data[:n],\n",
        "              final[:n].cpu()])\n",
        "            comparison = comparison.cpu()\n",
        "#             print(comparison.data)\n",
        "            save_image(com_img[:n].data,\n",
        "                       'compressed_' + str(epoch) +'.png', nrow=n)\n",
        "            save_image(comparison.data,\n",
        "                       'reconstruction_' + str(epoch) +'.png', nrow=n)\n",
        "\n",
        "  test_loss /= len(test_loader.dataset)\n",
        "  print('====> Test set loss: {:.4f}'.format(test_loss))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "sgm066wKMWVU",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Run the code"
      ]
    },
    {
      "metadata": {
        "id": "vwuk-i3jcnIQ",
        "colab_type": "code",
        "outputId": "c20016e1-917e-4161-cada-fc71c614572d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1136
        }
      },
      "cell_type": "code",
      "source": [
        "for epoch in range(1, EPOCHS+1):\n",
        "    train(epoch)\n",
        "    test(epoch)\n",
        "    if epoch == EPOCHS:\n",
        "      pass\n",
        "#         sample = Variable(torch.randn(64, args.hidden_size),16)\n",
        "#         sample = model.decode(sample)\n",
        "#         print(\"saving im\")\n",
        "#         save_image(sample.data.view(64, 3, 32, 32),\n",
        "# 'sample_' + str(epoch) + '.png')"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:52: UserWarning: size_average and reduce args will be deprecated, please use reduction='sum' instead.\n",
            "  warnings.warn(warning.format(ret))\n",
            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:10: 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",
            "  # Remove the CWD from sys.path while we load stuff.\n",
            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:16: 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",
            "  app.launch_new_instance()\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Train Epoch: 1 [0/50000 (0%)]\tLoss: 10073.152344\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "Process Process-1:\n",
            "Process Process-2:\n",
            "Traceback (most recent call last):\n",
            "Traceback (most recent call last):\n",
            "  File \"/usr/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
            "    self.run()\n",
            "  File \"/usr/lib/python3.6/multiprocessing/process.py\", line 258, in _bootstrap\n",
            "    self.run()\n",
            "  File \"/usr/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
            "    self._target(*self._args, **self._kwargs)\n",
            "  File \"/usr/lib/python3.6/multiprocessing/process.py\", line 93, in run\n",
            "    self._target(*self._args, **self._kwargs)\n",
            "  File \"/usr/local/lib/python3.6/dist-packages/torch/utils/data/dataloader.py\", line 96, in _worker_loop\n",
            "    r = index_queue.get(timeout=MANAGER_STATUS_CHECK_INTERVAL)\n",
            "  File \"/usr/local/lib/python3.6/dist-packages/torch/utils/data/dataloader.py\", line 96, in _worker_loop\n",
            "    r = index_queue.get(timeout=MANAGER_STATUS_CHECK_INTERVAL)\n",
            "  File \"/usr/lib/python3.6/multiprocessing/queues.py\", line 104, in get\n",
            "    if not self._poll(timeout):\n",
            "  File \"/usr/lib/python3.6/multiprocessing/queues.py\", line 104, in get\n",
            "    if not self._poll(timeout):\n",
            "  File \"/usr/lib/python3.6/multiprocessing/connection.py\", line 257, in poll\n",
            "    return self._poll(timeout)\n",
            "  File \"/usr/lib/python3.6/multiprocessing/connection.py\", line 257, in poll\n",
            "    return self._poll(timeout)\n",
            "  File \"/usr/lib/python3.6/multiprocessing/connection.py\", line 414, in _poll\n",
            "    r = wait([self], timeout)\n",
            "  File \"/usr/lib/python3.6/multiprocessing/connection.py\", line 414, in _poll\n",
            "    r = wait([self], timeout)\n",
            "  File \"/usr/lib/python3.6/multiprocessing/connection.py\", line 911, in wait\n",
            "    ready = selector.select(timeout)\n",
            "  File \"/usr/lib/python3.6/multiprocessing/connection.py\", line 911, in wait\n",
            "    ready = selector.select(timeout)\n",
            "  File \"/usr/lib/python3.6/selectors.py\", line 376, in select\n",
            "    fd_event_list = self._poll.poll(timeout)\n",
            "KeyboardInterrupt\n",
            "  File \"/usr/lib/python3.6/selectors.py\", line 376, in select\n",
            "    fd_event_list = self._poll.poll(timeout)\n",
            "KeyboardInterrupt\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-15-dbcc807c08e9>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mEPOCHS\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m     \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mepoch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m     \u001b[0mtest\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mepoch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mEPOCHS\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m       \u001b[0;32mpass\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-13-b8ece844b250>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(epoch)\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mVariable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m         \u001b[0mfinal\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresidual_img\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mupscaled_image\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcom_img\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0morig_im\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcuda\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      8\u001b[0m         \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfinal\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresidual_img\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mupscaled_image\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcom_img\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0morig_im\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    },
    {
      "metadata": {
        "id": "IdqFO814nFmd",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Reference : https://github.com/L1aoXingyu/pytorch-beginner/tree/master/08-AutoEncoder"
      ]
    },
    {
      "metadata": {
        "id": "lV2s0SMv2yO3",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#   model.save_state_dict()\n",
        "  torch.save(model.state_dict(), './net.pth')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "g0Jjq93jYN3C",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# from google.colab import drive\n",
        "# drive.mount('/content/gdrive') "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "BfAnSwModHG_",
        "colab_type": "code",
        "outputId": "db838e9b-2aa5-469a-9a47-24257fee4eaa",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 173
        }
      },
      "cell_type": "code",
      "source": [
        "model.load_state_dict(torch.load('net.pth'))\n",
        "\n",
        "def save_images():\n",
        "  epoch = EPOCHS\n",
        "  model.eval()\n",
        "  test_loss = 0\n",
        "  for i, (data, _) in enumerate(test_loader):\n",
        "        data = Variable(data, volatile=True)\n",
        "        final, residual_img, upscaled_image, com_img, orig_im = model(data.cuda())\n",
        "        test_loss += loss_function(final, residual_img, upscaled_image, com_img, orig_im).data[0]\n",
        "        if i == 3:\n",
        "#             save_image(final.data[0],'reconstruction_final',nrow=8)\n",
        "#             save_image(com_img.data[0],'com_img',nrow=8)\n",
        "            n = min(data.size(0), 6)\n",
        "            print(\"saving the image \"+str(n))\n",
        "            comparison = torch.cat([data[:n],\n",
        "              final[:n].cpu()])\n",
        "            comparison = comparison.cpu()\n",
        "#             print(comparison.data)\n",
        "            save_image(com_img[:1].data,\n",
        "                         'compressed_' + str(i) +'.png', nrow=n)\n",
        "            save_image(final[:1].data,\n",
        "                        'final_' + str(epoch) +'.png', nrow=n)\n",
        "            save_image(orig_im[:1].data,\n",
        "                        'original_' + str(epoch) +'.png', nrow=n)\n",
        "\n",
        "\n",
        "  test_loss /= len(test_loader.dataset)\n",
        "  print('====> Test set loss: {:.4f}'.format(test_loss))\n",
        "\n",
        "save_images()"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:8: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n",
            "  \n",
            "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py:52: UserWarning: size_average and reduce args will be deprecated, please use reduction='sum' instead.\n",
            "  warnings.warn(warning.format(ret))\n",
            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:10: 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",
            "  # Remove the CWD from sys.path while we load stuff.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "saving the image 6\n",
            "====> Test set loss: 176.2538\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "Q4jniD80b0gN",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}