autoencoder/denoising-autoencoder/Denoising_Autoencoder_Exercise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Denoising Autoencoder\n",
    "\n",
    "Sticking with the MNIST dataset, let's add noise to our data and see if we can define and train an autoencoder to _de_-noise the images.\n",
    "\n",
    "<img src='notebook_ims/autoencoder_denoise.png' width=70%/>\n",
    "\n",
    "Let's get started by importing our libraries and getting the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from torchvision import datasets\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "# convert data to torch.FloatTensor\n",
    "transform = transforms.ToTensor()\n",
    "\n",
    "# load the training and test datasets\n",
    "train_data = datasets.MNIST(root='data', train=True,\n",
    "                                   download=True, transform=transform)\n",
    "test_data = datasets.MNIST(root='data', train=False,\n",
    "                                  download=True, transform=transform)\n",
    "\n",
    "# Create training and test dataloaders\n",
    "num_workers = 0\n",
    "# how many samples per batch to load\n",
    "batch_size = 20\n",
    "\n",
    "# prepare data loaders\n",
    "train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, num_workers=num_workers)\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, num_workers=num_workers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "    \n",
    "# obtain one batch of training images\n",
    "dataiter = iter(train_loader)\n",
    "images, labels = dataiter.next()\n",
    "images = images.numpy()\n",
    "\n",
    "# get one image from the batch\n",
    "img = np.squeeze(images[0])\n",
    "\n",
    "fig = plt.figure(figsize = (5,5)) \n",
    "ax = fig.add_subplot(111)\n",
    "ax.imshow(img, cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Denoising\n",
    "\n",
    "As I've mentioned before, autoencoders like the ones you've built so far aren't too useful in practive. However, they can be used to denoise images quite successfully just by training the network on noisy images. We can create the noisy images ourselves by adding Gaussian noise to the training images, then clipping the values to be between 0 and 1.\n",
    "\n",
    ">**We'll use noisy images as input and the original, clean images as targets.** \n",
    "\n",
    "Below is an example of some of the noisy images I generated and the associated, denoised images.\n",
    "\n",
    "<img src='notebook_ims/denoising.png' />\n",
    "\n",
    "\n",
    "Since this is a harder problem for the network, we'll want to use _deeper_ convolutional layers here; layers with more feature maps. You might also consider adding additional layers. I suggest starting with a depth of 32 for the convolutional layers in the encoder, and the same depths going backward through the decoder.\n",
    "\n",
    "#### TODO: Build the network for the denoising autoencoder. Add deeper and/or additional layers compared to the model above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# define the NN architecture\n",
    "class ConvDenoiser(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ConvDenoiser, self).__init__()\n",
    "        ## encoder layers ##\n",
    "        \n",
    "        ## decoder layers ##\n",
    "        ## a kernel of 2 and a stride of 2 will increase the spatial dims by 2\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        ## encode ##\n",
    "        \n",
    "        ## decode ##\n",
    "                \n",
    "        return x\n",
    "\n",
    "# initialize the NN\n",
    "model = ConvDenoiser()\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Training\n",
    "\n",
    "We are only concerned with the training images, which we can get from the `train_loader`.\n",
    "\n",
    ">In this case, we are actually **adding some noise** to these images and we'll feed these `noisy_imgs` to our model. The model will produce reconstructed images based on the noisy input. But, we want it to produce _normal_ un-noisy images, and so, when we calculate the loss, we will still compare the reconstructed outputs to the original images!\n",
    "\n",
    "Because we're comparing pixel values in input and output images, it will be best to use a loss that is meant for a regression task. Regression is all about comparing quantities rather than probabilistic values. So, in this case, I'll use `MSELoss`. And compare output images and input images as follows:\n",
    "```\n",
    "loss = criterion(outputs, images)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# specify loss function\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "# specify loss function\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# number of epochs to train the model\n",
    "n_epochs = 20\n",
    "\n",
    "# for adding noise to images\n",
    "noise_factor=0.5\n",
    "\n",
    "for epoch in range(1, n_epochs+1):\n",
    "    # monitor training loss\n",
    "    train_loss = 0.0\n",
    "    \n",
    "    ###################\n",
    "    # train the model #\n",
    "    ###################\n",
    "    for data in train_loader:\n",
    "        # _ stands in for labels, here\n",
    "        # no need to flatten images\n",
    "        images, _ = data\n",
    "        \n",
    "        ## add random noise to the input images\n",
    "        noisy_imgs = images + noise_factor * torch.randn(*images.shape)\n",
    "        # Clip the images to be between 0 and 1\n",
    "        noisy_imgs = np.clip(noisy_imgs, 0., 1.)\n",
    "                \n",
    "        # clear the gradients of all optimized variables\n",
    "        optimizer.zero_grad()\n",
    "        ## forward pass: compute predicted outputs by passing *noisy* images to the model\n",
    "        outputs = model(noisy_imgs)\n",
    "        # calculate the loss\n",
    "        # the \"target\" is still the original, not-noisy images\n",
    "        loss = criterion(outputs, images)\n",
    "        # backward pass: compute gradient of the loss with respect to model parameters\n",
    "        loss.backward()\n",
    "        # perform a single optimization step (parameter update)\n",
    "        optimizer.step()\n",
    "        # update running training loss\n",
    "        train_loss += loss.item()*images.size(0)\n",
    "            \n",
    "    # print avg training statistics \n",
    "    train_loss = train_loss/len(train_loader)\n",
    "    print('Epoch: {} \\tTraining Loss: {:.6f}'.format(\n",
    "        epoch, \n",
    "        train_loss\n",
    "        ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checking out the results\n",
    "\n",
    "Here I'm adding noise to the test images and passing them through the autoencoder. It does a suprising great job of removing the noise, even though it's sometimes difficult to tell what the original number is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# obtain one batch of test images\n",
    "dataiter = iter(test_loader)\n",
    "images, labels = dataiter.next()\n",
    "\n",
    "# add noise to the test images\n",
    "noisy_imgs = images + noise_factor * torch.randn(*images.shape)\n",
    "noisy_imgs = np.clip(noisy_imgs, 0., 1.)\n",
    "\n",
    "# get sample outputs\n",
    "output = model(noisy_imgs)\n",
    "# prep images for display\n",
    "noisy_imgs = noisy_imgs.numpy()\n",
    "\n",
    "# output is resized into a batch of iages\n",
    "output = output.view(batch_size, 1, 28, 28)\n",
    "# use detach when it's an output that requires_grad\n",
    "output = output.detach().numpy()\n",
    "\n",
    "# plot the first ten input images and then reconstructed images\n",
    "fig, axes = plt.subplots(nrows=2, ncols=10, sharex=True, sharey=True, figsize=(25,4))\n",
    "\n",
    "# input images on top row, reconstructions on bottom\n",
    "for noisy_imgs, row in zip([noisy_imgs, output], axes):\n",
    "    for img, ax in zip(noisy_imgs, row):\n",
    "        ax.imshow(np.squeeze(img), cmap='gray')\n",
    "        ax.get_xaxis().set_visible(False)\n",
    "        ax.get_yaxis().set_visible(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "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": 2
}