gan.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cuda'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import pandas as pd\n",
    "import time\n",
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Variable\n",
    "import torch.nn.utils.rnn as rnn\n",
    "from torch.utils.data import Dataset, DataLoader, BatchSampler, RandomSampler\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "DEVICE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "male = np.load(\"male_spect.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10000,)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "male.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(348, 129)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "male[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainDataset(Dataset):\n",
    "    def __init__(self, train):\n",
    "        self.train = train\n",
    "        lens = [a.shape[0] for a in train]\n",
    "        self.max_len = np.max(lens)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.train.shape[0]\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        diff = self.max_len - self.train[idx].shape[0]\n",
    "        train = np.pad(self.train[idx], [(0, diff),(0,0)], \"constant\", constant_values=0)\n",
    "        return train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = TrainDataset(male)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5328, 129)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset[4].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 20\n",
    "train_loader = DataLoader(train_dataset,\n",
    "                        batch_sampler=BatchSampler(RandomSampler(train_dataset), batch_size, False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([20, 5328, 129])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "iter(train_loader).next().shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def weights_init_normal(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Conv') != -1:\n",
    "        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
    "    elif classname.find('BatchNorm2d') != -1:\n",
    "        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
    "        torch.nn.init.constant_(m.bias.data, 0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self, input_shape):\n",
    "        super(Generator, self).__init__()\n",
    "\n",
    "        self.init_size = input_shape[1] // 4\n",
    "        self.l1 = nn.Sequential(nn.Linear(input_shape[0], 128*self.init_size**2))\n",
    "\n",
    "        self.conv_blocks = nn.Sequential(\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.Upsample(scale_factor=2),\n",
    "            nn.functional.interpolate(scale_factor=2),\n",
    "            nn.Conv2d(128, 128, 3, stride=1, padding=1),\n",
    "            nn.BatchNorm2d(128, 0.8),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Upsample(scale_factor=2),\n",
    "            nn.Conv2d(128, 64, 3, stride=1, padding=1),\n",
    "            nn.BatchNorm2d(64, 0.8),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(64, 1, 3, stride=1, padding=1),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "\n",
    "    def forward(self, z):\n",
    "        out = self.l1(z)\n",
    "        out = out.view(out.shape[0], 128, self.init_size, self.init_size)\n",
    "        img = self.conv_blocks(out)\n",
    "        return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/jet/var/python/lib/python3.6/site-packages/torch/nn/modules/upsampling.py:122: UserWarning: nn.Upsampling is deprecated. Use nn.functional.interpolate instead.\n",
      "  warnings.warn(\"nn.Upsampling is deprecated. Use nn.functional.interpolate instead.\")\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "$ Torch: not enough memory: you tried to allocate 4GB. Buy new RAM! at /jet/tmp/build/aten/src/TH/THGeneral.cpp:204",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-44-352954702ee2>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;31m# Generate a batch of images\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m \u001b[0mgen_imgs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgenerator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mz\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/var/python/lib/python3.6/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    475\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    476\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 477\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    478\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\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[1;32m    479\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-42-58d111097bbf>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, z)\u001b[0m\n\u001b[1;32m     23\u001b[0m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0ml1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mz\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m128\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minit_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minit_size\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 25\u001b[0;31m         \u001b[0mimg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconv_blocks\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     26\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mimg\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/var/python/lib/python3.6/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    475\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    476\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 477\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    478\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\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[1;32m    479\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/var/python/lib/python3.6/site-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m     89\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     90\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\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[0;32m---> 91\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     92\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     93\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/var/python/lib/python3.6/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    475\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    476\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 477\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    478\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\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[1;32m    479\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/var/python/lib/python3.6/site-packages/torch/nn/modules/conv.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    299\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    300\u001b[0m         return F.conv2d(input, self.weight, self.bias, self.stride,\n\u001b[0;32m--> 301\u001b[0;31m                         self.padding, self.dilation, self.groups)\n\u001b[0m\u001b[1;32m    302\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    303\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: $ Torch: not enough memory: you tried to allocate 4GB. Buy new RAM! at /jet/tmp/build/aten/src/TH/THGeneral.cpp:204"
     ]
    }
   ],
   "source": [
    "generator = Generator(input_shape= (5328, 129))\n",
    "generator.cuda()\n",
    "z = torch.Tensor(np.random.normal(0, 1, (129, 5328)))\n",
    "z.cuda()\n",
    "\n",
    "# Generate a batch of images\n",
    "gen_imgs = generator(z)\n",
    "\n",
    "print(gen_imgs.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def weights_init_normal(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Conv') != -1:\n",
    "        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
    "    elif classname.find('BatchNorm2d') != -1:\n",
    "        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
    "        torch.nn.init.constant_(m.bias.data, 0.0)\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self, input_shape):\n",
    "        super(Generator, self).__init__()\n",
    "\n",
    "        self.init_size = opt.img_size // 4\n",
    "        self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128*self.init_size**2))\n",
    "\n",
    "        self.conv_blocks = nn.Sequential(\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.Upsample(scale_factor=2),\n",
    "            nn.Conv2d(128, 128, 3, stride=1, padding=1),\n",
    "            nn.BatchNorm2d(128, 0.8),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Upsample(scale_factor=2),\n",
    "            nn.Conv2d(128, 64, 3, stride=1, padding=1),\n",
    "            nn.BatchNorm2d(64, 0.8),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "\n",
    "    def forward(self, z):\n",
    "        out = self.l1(z)\n",
    "        out = out.view(out.shape[0], 128, self.init_size, self.init_size)\n",
    "        img = self.conv_blocks(out)\n",
    "        return img\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Discriminator, self).__init__()\n",
    "\n",
    "        def discriminator_block(in_filters, out_filters, bn=True):\n",
    "            block = [   nn.Conv2d(in_filters, out_filters, 3, 2, 1),\n",
    "                        nn.LeakyReLU(0.2, inplace=True),\n",
    "                        nn.Dropout2d(0.25)]\n",
    "            if bn:\n",
    "                block.append(nn.BatchNorm2d(out_filters, 0.8))\n",
    "            return block\n",
    "\n",
    "        self.model = nn.Sequential(\n",
    "            *discriminator_block(opt.channels, 16, bn=False),\n",
    "            *discriminator_block(16, 32),\n",
    "            *discriminator_block(32, 64),\n",
    "            *discriminator_block(64, 128),\n",
    "        )\n",
    "\n",
    "        # The height and width of downsampled image\n",
    "        ds_size = opt.img_size // 2**4\n",
    "        self.adv_layer = nn.Sequential( nn.Linear(128*ds_size**2, 1),\n",
    "                                        nn.Sigmoid())\n",
    "\n",
    "    def forward(self, img):\n",
    "        out = self.model(img)\n",
    "        out = out.view(out.shape[0], -1)\n",
    "        validity = self.adv_layer(out)\n",
    "\n",
    "        return validity\n",
    "\n",
    "# Loss function\n",
    "adversarial_loss = torch.nn.BCELoss()\n",
    "\n",
    "# Initialize generator and discriminator\n",
    "generator = Generator()\n",
    "discriminator = Discriminator()\n",
    "\n",
    "if cuda:\n",
    "    generator.cuda()\n",
    "    discriminator.cuda()\n",
    "    adversarial_loss.cuda()\n",
    "\n",
    "# Initialize weights\n",
    "generator.apply(weights_init_normal)\n",
    "discriminator.apply(weights_init_normal)\n",
    "\n",
    "# Configure data loader\n",
    "os.makedirs('../../data/mnist', exist_ok=True)\n",
    "dataloader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST('../../data/mnist', train=True, download=True,\n",
    "                   transform=transforms.Compose([\n",
    "                       transforms.Resize(opt.img_size),\n",
    "                       transforms.ToTensor(),\n",
    "                       transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "                   ])),\n",
    "    batch_size=opt.batch_size, shuffle=True)\n",
    "\n",
    "# Optimizers\n",
    "optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n",
    "optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))\n",
    "\n",
    "Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor\n",
    "\n",
    "# ----------\n",
    "#  Training\n",
    "# ----------\n",
    "\n",
    "for epoch in range(opt.n_epochs):\n",
    "    for i, (imgs, _) in enumerate(dataloader):\n",
    "\n",
    "        # Adversarial ground truths\n",
    "        valid = Variable(Tensor(imgs.shape[0], 1).fill_(1.0), requires_grad=False)\n",
    "        fake = Variable(Tensor(imgs.shape[0], 1).fill_(0.0), requires_grad=False)\n",
    "\n",
    "        # Configure input\n",
    "        real_imgs = Variable(imgs.type(Tensor))\n",
    "\n",
    "        # -----------------\n",
    "        #  Train Generator\n",
    "        # -----------------\n",
    "\n",
    "        optimizer_G.zero_grad()\n",
    "\n",
    "        # Sample noise as generator input\n",
    "        z = Variable(Tensor(np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))))\n",
    "\n",
    "        # Generate a batch of images\n",
    "        gen_imgs = generator(z)\n",
    "\n",
    "        # Loss measures generator's ability to fool the discriminator\n",
    "        g_loss = adversarial_loss(discriminator(gen_imgs), valid)\n",
    "\n",
    "        g_loss.backward()\n",
    "        optimizer_G.step()\n",
    "\n",
    "        # ---------------------\n",
    "        #  Train Discriminator\n",
    "        # ---------------------\n",
    "\n",
    "        optimizer_D.zero_grad()\n",
    "\n",
    "        # Measure discriminator's ability to classify real from generated samples\n",
    "        real_loss = adversarial_loss(discriminator(real_imgs), valid)\n",
    "        fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)\n",
    "        d_loss = (real_loss + fake_loss) / 2\n",
    "\n",
    "        d_loss.backward()\n",
    "        optimizer_D.step()\n",
    "\n",
    "        print (\"[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]\" % (epoch, opt.n_epochs, i, len(dataloader),\n",
    "                                                            d_loss.item(), g_loss.item()))\n",
    "\n",
    "        batches_done = epoch * len(dataloader) + i\n",
    "        if batches_done % opt.sample_interval == 0:\n",
    "            save_image(gen_imgs.data[:25], 'images/%d.png' % batches_done, nrow=5, normalize=True)"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}