Image_Generative_model/Denoising_AE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.utils import save_image\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import os\n",
    "import glob\n",
    "import PIL\n",
    "from PIL import Image\n",
    "from scipy.misc import toimage\n",
    "from torch.utils import data as D\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def denorm(x):\n",
    "    out = (x + 1) / 2\n",
    "    return out.clamp(0, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "img_size = (28, 28, 1) \n",
    "hidden_size = 256\n",
    "\n",
    "num_epochs = 100\n",
    "batch_size = 64\n",
    "\n",
    "dataset_dir = \"./data/FMNIST\"\n",
    "sample_dir = \"./result_denoising_ae_FMNIST/\" \n",
    "\n",
    "if not os.path.exists(dataset_dir):\n",
    "    os.makedirs(dataset_dir)\n",
    "\n",
    "if not os.path.exists(sample_dir):\n",
    "    os.makedirs(sample_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "            transforms.Resize(img_size[0]),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(mean=(0.5, 0.5, 0.5),std=(0.5, 0.5, 0.5))])\n",
    "    \n",
    "trainset = torchvision.datasets.FashionMNIST(root=dataset_dir, train=True ,transform=transform, download=True)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class Denoising_autoencoder(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size):\n",
    "        super(Denoising_autoencoder, self).__init__()\n",
    "        \n",
    "        self.network = nn.Sequential(\n",
    "            # (-1, 1, 28, 28) --> (-1, 256, 14, 14)\n",
    "            nn.Conv2d(input_size[2], hidden_size, kernel_size=3, padding=1, stride=2, bias=False),\n",
    "            nn.BatchNorm2d(hidden_size),\n",
    "            nn.ReLU(True),\n",
    "            \n",
    "            # (-1, 256, 14, 14) --> (-1, 1024, 7, 7)\n",
    "            nn.Conv2d(hidden_size, hidden_size*4, kernel_size=3, padding=1, stride=2, bias=False),\n",
    "            nn.BatchNorm2d(hidden_size*4),\n",
    "            nn.ReLU(True),\n",
    "            \n",
    "            # (-1, 1024, 7, 7) --> (-1, 2048, 7, 7)\n",
    "            nn.Conv2d(hidden_size*4, hidden_size*8, kernel_size=3, padding=1, stride=1, bias=False),\n",
    "            nn.BatchNorm2d(hidden_size*8),\n",
    "            nn.ReLU(True),\n",
    "            \n",
    "            # (-1, 2048, 7, 7) --> (-1, 1024, 7, 7)\n",
    "            nn.Conv2d(hidden_size*8, hidden_size*4, kernel_size=3, padding=1, stride=1, bias=False),\n",
    "            nn.BatchNorm2d(hidden_size*4),\n",
    "            nn.ReLU(True),\n",
    "            \n",
    "            # (-1, 1024, 7, 7) --> (-1, 256, 14, 14)\n",
    "            nn.ConvTranspose2d(hidden_size*4, hidden_size, kernel_size=3, padding=1, stride=2, output_padding=1, bias=False),\n",
    "            nn.BatchNorm2d(hidden_size),\n",
    "            nn.ReLU(True),\n",
    "            \n",
    "            # (-1, 256, 14, 14) --> (-1, 1, 28, 28)\n",
    "            nn.ConvTranspose2d(hidden_size, input_size[2], kernel_size=3, padding=1, stride=2, output_padding=1, bias=False),\n",
    "            nn.BatchNorm2d(input_size[2]),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        output = self.network(x)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda:0\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Denoising_autoencoder(\n",
       "  (network): Sequential(\n",
       "    (0): Conv2d(1, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU(inplace)\n",
       "    (3): Conv2d(256, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (4): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (5): ReLU(inplace)\n",
       "    (6): Conv2d(1024, 2048, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (7): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (8): ReLU(inplace)\n",
       "    (9): Conv2d(2048, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (10): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (11): ReLU(inplace)\n",
       "    (12): ConvTranspose2d(1024, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1), bias=False)\n",
       "    (13): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (14): ReLU(inplace)\n",
       "    (15): ConvTranspose2d(256, 1, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1), bias=False)\n",
       "    (16): BatchNorm2d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (17): Tanh()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "AE_model = Denoising_autoencoder(img_size, hidden_size) \n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)\n",
    "\n",
    "AE_model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "criterion = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(AE_model.parameters(), lr=0.01)\n",
    "#optimizer = torch.optim.SGD(AE_model.parameters(), lr=0.1, momentum=0.9)\n",
    "init_lr = 0.1\n",
    "\n",
    "mean = 0\n",
    "std = 0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [0/100], Step [800/938], mse loss: 0.0082, lr: 0.01000\n",
      "Epoch [1/100], Step [800/938], mse loss: 0.0067, lr: 0.01000\n",
      "Epoch [2/100], Step [800/938], mse loss: 0.0076, lr: 0.01000\n",
      "Epoch [3/100], Step [800/938], mse loss: 0.0067, lr: 0.01000\n",
      "Epoch [4/100], Step [800/938], mse loss: 0.0075, lr: 0.01000\n",
      "Epoch [5/100], Step [800/938], mse loss: 0.0071, lr: 0.01000\n",
      "Epoch [6/100], Step [800/938], mse loss: 0.0067, lr: 0.01000\n",
      "Epoch [7/100], Step [800/938], mse loss: 0.0067, lr: 0.01000\n",
      "Epoch [8/100], Step [800/938], mse loss: 0.0069, lr: 0.01000\n",
      "Epoch [9/100], Step [800/938], mse loss: 0.0072, lr: 0.01000\n",
      "Epoch [10/100], Step [800/938], mse loss: 0.0067, lr: 0.01000\n",
      "Epoch [11/100], Step [800/938], mse loss: 0.0060, lr: 0.01000\n",
      "Epoch [12/100], Step [800/938], mse loss: 0.0068, lr: 0.01000\n",
      "Epoch [13/100], Step [800/938], mse loss: 0.0063, lr: 0.01000\n",
      "Epoch [14/100], Step [800/938], mse loss: 0.0090, lr: 0.01000\n",
      "Epoch [15/100], Step [800/938], mse loss: 0.0060, lr: 0.01000\n",
      "Epoch [16/100], Step [800/938], mse loss: 0.0075, lr: 0.01000\n",
      "Epoch [17/100], Step [800/938], mse loss: 0.0065, lr: 0.01000\n",
      "Epoch [18/100], Step [800/938], mse loss: 0.0065, lr: 0.01000\n",
      "Epoch [19/100], Step [800/938], mse loss: 0.0062, lr: 0.01000\n",
      "Epoch [20/100], Step [800/938], mse loss: 0.0062, lr: 0.00100\n",
      "Epoch [21/100], Step [800/938], mse loss: 0.0061, lr: 0.00100\n",
      "Epoch [22/100], Step [800/938], mse loss: 0.0057, lr: 0.00100\n",
      "Epoch [23/100], Step [800/938], mse loss: 0.0058, lr: 0.00100\n",
      "Epoch [24/100], Step [800/938], mse loss: 0.0059, lr: 0.00100\n",
      "Epoch [25/100], Step [800/938], mse loss: 0.0062, lr: 0.00100\n",
      "Epoch [26/100], Step [800/938], mse loss: 0.0062, lr: 0.00100\n",
      "Epoch [27/100], Step [800/938], mse loss: 0.0056, lr: 0.00100\n",
      "Epoch [28/100], Step [800/938], mse loss: 0.0062, lr: 0.00100\n",
      "Epoch [29/100], Step [800/938], mse loss: 0.0062, lr: 0.00100\n",
      "Epoch [30/100], Step [800/938], mse loss: 0.0060, lr: 0.00100\n",
      "Epoch [31/100], Step [800/938], mse loss: 0.0055, lr: 0.00100\n",
      "Epoch [32/100], Step [800/938], mse loss: 0.0058, lr: 0.00100\n",
      "Epoch [33/100], Step [800/938], mse loss: 0.0055, lr: 0.00100\n",
      "Epoch [34/100], Step [800/938], mse loss: 0.0059, lr: 0.00100\n",
      "Epoch [35/100], Step [800/938], mse loss: 0.0053, lr: 0.00100\n",
      "Epoch [36/100], Step [800/938], mse loss: 0.0062, lr: 0.00100\n",
      "Epoch [37/100], Step [800/938], mse loss: 0.0056, lr: 0.00100\n",
      "Epoch [38/100], Step [800/938], mse loss: 0.0056, lr: 0.00100\n",
      "Epoch [39/100], Step [800/938], mse loss: 0.0064, lr: 0.00100\n",
      "Epoch [40/100], Step [800/938], mse loss: 0.0057, lr: 0.00010\n",
      "Epoch [41/100], Step [800/938], mse loss: 0.0061, lr: 0.00010\n",
      "Epoch [42/100], Step [800/938], mse loss: 0.0062, lr: 0.00010\n",
      "Epoch [43/100], Step [800/938], mse loss: 0.0061, lr: 0.00010\n",
      "Epoch [44/100], Step [800/938], mse loss: 0.0066, lr: 0.00010\n",
      "Epoch [45/100], Step [800/938], mse loss: 0.0065, lr: 0.00010\n",
      "Epoch [46/100], Step [800/938], mse loss: 0.0058, lr: 0.00010\n",
      "Epoch [47/100], Step [800/938], mse loss: 0.0056, lr: 0.00010\n",
      "Epoch [48/100], Step [800/938], mse loss: 0.0057, lr: 0.00010\n",
      "Epoch [49/100], Step [800/938], mse loss: 0.0064, lr: 0.00010\n",
      "Epoch [50/100], Step [800/938], mse loss: 0.0059, lr: 0.00010\n",
      "Epoch [51/100], Step [800/938], mse loss: 0.0062, lr: 0.00010\n",
      "Epoch [52/100], Step [800/938], mse loss: 0.0059, lr: 0.00010\n",
      "Epoch [53/100], Step [800/938], mse loss: 0.0065, lr: 0.00010\n",
      "Epoch [54/100], Step [800/938], mse loss: 0.0056, lr: 0.00010\n",
      "Epoch [55/100], Step [800/938], mse loss: 0.0056, lr: 0.00010\n",
      "Epoch [56/100], Step [800/938], mse loss: 0.0056, lr: 0.00010\n",
      "Epoch [57/100], Step [800/938], mse loss: 0.0060, lr: 0.00010\n",
      "Epoch [58/100], Step [800/938], mse loss: 0.0057, lr: 0.00010\n",
      "Epoch [59/100], Step [800/938], mse loss: 0.0061, lr: 0.00010\n",
      "Epoch [60/100], Step [800/938], mse loss: 0.0067, lr: 0.00001\n",
      "Epoch [61/100], Step [800/938], mse loss: 0.0061, lr: 0.00001\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-33-85c8fe882478>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m     18\u001b[0m         \u001b[0moutputs\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mAE_model\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mimages_noise\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     19\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 20\u001b[1;33m         \u001b[0mmse_loss\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0moutputs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mimages\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     21\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     22\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mC:\\Users\\Lee.hoseong\\Anaconda3\\lib\\site-packages\\torch\\nn\\modules\\module.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, *input, **kwargs)\u001b[0m\n\u001b[0;32m    489\u001b[0m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    490\u001b[0m         \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 491\u001b[1;33m             \u001b[0mresult\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    492\u001b[0m         \u001b[1;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    493\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mC:\\Users\\Lee.hoseong\\Anaconda3\\lib\\site-packages\\torch\\nn\\modules\\loss.py\u001b[0m in \u001b[0;36mforward\u001b[1;34m(self, input, target)\u001b[0m\n\u001b[0;32m    370\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtarget\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    371\u001b[0m         \u001b[0m_assert_no_grad\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtarget\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 372\u001b[1;33m         \u001b[1;32mreturn\u001b[0m \u001b[0mF\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmse_loss\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtarget\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0msize_average\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msize_average\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mreduce\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mreduce\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    373\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    374\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mC:\\Users\\Lee.hoseong\\Anaconda3\\lib\\site-packages\\torch\\nn\\functional.py\u001b[0m in \u001b[0;36mmse_loss\u001b[1;34m(input, target, size_average, reduce)\u001b[0m\n\u001b[0;32m   1567\u001b[0m     \"\"\"\n\u001b[0;32m   1568\u001b[0m     return _pointwise_loss(lambda a, b: (a - b) ** 2, torch._C._nn.mse_loss,\n\u001b[1;32m-> 1569\u001b[1;33m                            input, target, size_average, reduce)\n\u001b[0m\u001b[0;32m   1570\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1571\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mC:\\Users\\Lee.hoseong\\Anaconda3\\lib\\site-packages\\torch\\nn\\functional.py\u001b[0m in \u001b[0;36m_pointwise_loss\u001b[1;34m(lambd, lambd_optimized, input, target, size_average, reduce)\u001b[0m\n\u001b[0;32m   1535\u001b[0m         \u001b[1;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0md\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0msize_average\u001b[0m \u001b[1;32melse\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msum\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0md\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1536\u001b[0m     \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m-> 1537\u001b[1;33m         \u001b[1;32mreturn\u001b[0m \u001b[0mlambd_optimized\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0minput\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtarget\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0msize_average\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mreduce\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m   1538\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m   1539\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# Start training\n",
    "total_step = len(trainloader)\n",
    "for epoch in range(num_epochs):\n",
    "    if epoch % 20 == 0:\n",
    "        optimizer = torch.optim.Adam(AE_model.parameters(), lr=init_lr)\n",
    "        #optimizer = torch.optim.SGD(AE_model.parameters(), lr=init_lr, momentum=0.9)\n",
    "\n",
    "        init_lr /= 10\n",
    "    \n",
    "    for i, (images, _) in enumerate(trainloader):\n",
    "        # add noise\n",
    "       \n",
    "        noise = torch.tensor(images.data.new(images.size()).normal_(mean, std))\n",
    "        images_noise = images + noise\n",
    "        images = images.to(device)\n",
    "        images_noise = images_noise.to(device)\n",
    "\n",
    "        outputs = AE_model(images_noise)\n",
    "        \n",
    "        mse_loss = criterion(outputs, images)\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        mse_loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        if (i+1) % 800 == 0:\n",
    "            print('Epoch [{}/{}], Step [{}/{}], mse loss: {:.4f}, lr: {:.5f}' \n",
    "                  .format(epoch, num_epochs, i+1, total_step, mse_loss.item(), init_lr))\n",
    "    \n",
    "    # Save real images and output_images\n",
    "    images_noise = images_noise.reshape(images_noise.size()[0], 1, 28, 28)\n",
    "    output_images = outputs.reshape(outputs.size()[0], 1, 28, 28)\n",
    "        \n",
    "    results = torch.cat([images_noise, output_images], dim=2)\n",
    "    \n",
    "    save_image(denorm(results), os.path.join(sample_dir, 'ae_model_result-{}.png'.format(epoch+1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [Anaconda3]",
   "language": "python",
   "name": "Python [Anaconda3]"
  },
  "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}