End_to_end_image_compression_stl10.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "End_to_end_image_compression_stl10.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_stl10.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": "58aed18b-10fb-4c00-8f04-c7d3582973ca",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 306
        }
      },
      "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",
        "\n",
        "# Install and import torch\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",
        "\n",
        "#install and import pillow\n",
        "!pip install Pillow==4.0.0\n",
        "!pip install PIL\n",
        "!pip install image\n",
        "from PIL import Image"
      ],
      "execution_count": 1,
      "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: django in /usr/local/lib/python3.6/dist-packages (from image) (2.1.4)\n",
            "Requirement already satisfied: pillow in /usr/local/lib/python3.6/dist-packages (from image) (4.0.0)\n",
            "Requirement already satisfied: pytz in /usr/local/lib/python3.6/dist-packages (from django->image) (2018.7)\n",
            "Requirement already satisfied: olefile in /usr/local/lib/python3.6/dist-packages (from pillow->image) (0.46)\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": "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 = 96\n",
        "WIDTH = 96\n",
        "EPOCHS = 200\n",
        "LOG_INTERVAL = 50\n",
        "BATCH_SIZE = 16"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "I6xA8tQb1ckr",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Load Data"
      ]
    },
    {
      "metadata": {
        "id": "CtMMu7Dq1TIB",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "7d3227d7-8cab-4788-b081-6c4f64e86166"
      },
      "cell_type": "code",
      "source": [
        "img_transform = transforms.Compose([\n",
        "    transforms.ToTensor()\n",
        "])\n",
        "\n",
        "trainset = datasets.STL10(root='./data', split='train',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.STL10(root='./data', split='test',download=True, transform=img_transform)\n",
        "\n",
        "test_loader = torch.utils.data.DataLoader(testset, batch_size=16,shuffle=False, num_workers=2)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Files already downloaded and verified\n",
            "Files already downloaded and verified\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "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(5):\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": "RrnHrzxRIP0q",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "4f3256aa-29c9-4c34-b794-7fc5d5bcd7da"
      },
      "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",
        "print(\"GPU available ? \"+str(CUDA))\n",
        "  \n",
        "optimizer = optim.Adam(model.parameters(), lr=1e-3)"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "GPU available ? True\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "VHy3FuhX1oo7",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Loss function"
      ]
    },
    {
      "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": "IyZnGma91rwV",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### define Train and Test methods"
      ]
    },
    {
      "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",
        "        if CUDA:\n",
        "          final, residual_img, upscaled_image, com_img, orig_im = model(data.cuda())\n",
        "        else :\n",
        "          final, residual_img, upscaled_image, com_img, orig_im = model(data)\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",
        "        if CUDA:\n",
        "          final, residual_img, upscaled_image, com_img, orig_im = model(data.cuda())\n",
        "        else :\n",
        "          final, residual_img, upscaled_image, com_img, orig_im = model(data)\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",
        "\n",
        "\n",
        "  test_loss /= len(test_loader.dataset)\n",
        "  print('====> Test set loss: {:.4f}'.format(test_loss))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "_hiDc-Zh100e",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Program Execution"
      ]
    },
    {
      "metadata": {
        "id": "vwuk-i3jcnIQ",
        "colab_type": "code",
        "outputId": "2703de84-b2ec-4fa8-ab9d-088540be5604",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 394
        }
      },
      "cell_type": "code",
      "source": [
        "for epoch in range(1, EPOCHS+1):\n",
        "    print(\"Epoch : \"+str(epoch))\n",
        "    train(epoch)\n",
        "    test(epoch)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch : 1\n"
          ],
          "name": "stdout"
        },
        {
          "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:13: 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",
            "  del sys.path[0]\n",
            "/usr/local/lib/python3.6/dist-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": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Train Epoch: 1 [0/5000 (0%)]\tLoss: 33825.121094\n",
            "Train Epoch: 1 [800/5000 (16%)]\tLoss: 5496.435059\n",
            "Train Epoch: 1 [1600/5000 (32%)]\tLoss: 3329.267578\n",
            "Train Epoch: 1 [2400/5000 (48%)]\tLoss: 846.066956\n",
            "Train Epoch: 1 [3200/5000 (64%)]\tLoss: 622.065430\n",
            "Train Epoch: 1 [4000/5000 (80%)]\tLoss: 744.986450\n",
            "Train Epoch: 1 [4800/5000 (96%)]\tLoss: 495.083740\n",
            "====> Epoch: 1 Average loss: 2918.1763\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:6: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n",
            "  \n",
            "/usr/local/lib/python3.6/dist-packages/ipykernel_launcher.py:11: 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",
            "  # This is added back by InteractiveShellApp.init_path()\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "====> Test set loss: 837.0466\n",
            "Epoch : 2\n",
            "Train Epoch: 2 [0/5000 (0%)]\tLoss: 632.405334\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "IdqFO814nFmd",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Reference : https://github.com/L1aoXingyu/pytorch-beginner/tree/master/08-AutoEncoder"
      ]
    },
    {
      "metadata": {
        "id": "g0Jjq93jYN3C",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# from google.colab import drive\n",
        "# drive.mount('/content/gdrive') \n",
        "torch.save(model.state_dict(), './net.pth')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "46oJA6aMirJB",
        "colab_type": "code",
        "colab": {}
      },
      "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",
        "        if CUDA:\n",
        "          final, residual_img, upscaled_image, com_img, orig_im = model(data.cuda())\n",
        "        else:\n",
        "          final, residual_img, upscaled_image, com_img, orig_im = model(data)\n",
        "\n",
        "        test_loss += loss_function(final, residual_img, upscaled_image, com_img, orig_im).data[0]\n",
        "        if True:\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[i].data,\n",
        "                         'compressed_' + str(i) +'.png', nrow=n)\n",
        "            save_image(final[i].data,\n",
        "                        'final_' + str(i) +'.png', nrow=n)\n",
        "            save_image(orig_im[i].data,\n",
        "                        'original_' + str(i) +'.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": []
    }
  ]
}