{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"from __future__ import print_function\n#%matplotlib inline\nimport argparse\nimport os\nimport random\nimport torch\nimport torch.nn as nn\nimport torch.nn.parallel\nimport torch.backends.cudnn as cudnn\nimport torch.optim as optim\nimport torch.utils.data\nimport torchvision.datasets as dset\nimport torchvision.transforms as transforms\nimport torchvision.utils as vutils\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport matplotlib.animation as animation\nfrom IPython.display import HTML\nfrom torchvision.utils import save_image\nfrom torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau, CosineAnnealingLR","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\nseed_everything()\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Check Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"len(os.listdir('../input/all-dogs/all-dogs'))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Set Parameters"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Setting parameters\ndataroot = '../input/all-dogs'\nworkers = 2\n\nbatch_size = 128\nimage_size = 64\n\n# Number of channels\nnc = 3\n# Latent vector (i.e. size of generator input)\nnz = 100\n# Size of feature maps in generator\nngf = 64\n# Size of feature maps in discriminator\nndf = 64\n\n# Number of training epochs\nnum_epochs = 100\n# Learning rate for optimizers\nlr = 0.0005\n# Beta1 hyperparam for Adam optimizers\nbeta1 = 0.5\nngpu = 2","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Preparing data"},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = dset.ImageFolder(root=dataroot,\n                           transform=transforms.Compose([\n                               transforms.Resize(image_size),\n                               transforms.RandomHorizontalFlip(),\n                               transforms.CenterCrop(image_size),\n                               transforms.ToTensor(),\n                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n                           ]))\n\n\n# Create the dataloader\ndataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,\n                                         shuffle=True, num_workers=workers)\n\n# Decide which device we want to run on\ndevice = torch.device(\"cuda\" if (torch.cuda.is_available() and ngpu > 0) else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"real_batch = next(iter(dataloader))\n\nplt.figure(figsize=(10,10))\nplt.axis(\"off\")\nplt.title(\"Training Images\");\nplt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], \n                                         padding=2, normalize=True).cpu(),(1,2,0)));","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{},"cell_type":"markdown","source":"## Set Weight"},{"metadata":{"trusted":true},"cell_type":"code","source":"def weights_init(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:\n        nn.init.normal_(m.weight.data, 0.0, 0.02)\n    elif classname.find('BatchNorm') != -1:\n        nn.init.normal_(m.weight.data, 1.0, 0.02)\n        nn.init.constant_(m.bias.data, 0)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## DCGAN Architecture"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, ngpu):\n        super(Generator, self).__init__()\n        self.ngpu = ngpu\n        self.main = nn.Sequential(\n            # input is Z, going into a convolution\n            nn.ConvTranspose2d( nz, ngf * 8, 4, 1, 0, bias=False),\n            nn.BatchNorm2d(ngf * 8),\n            nn.ReLU(True),\n            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ngf * 4),\n            nn.ReLU(True),\n            nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ngf * 2),\n            nn.ReLU(True),\n            nn.ConvTranspose2d( ngf * 2, ngf, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ngf),\n            nn.ReLU(True),\n            nn.ConvTranspose2d( ngf, nc, 4, 2, 1, bias=False),\n            nn.Tanh()\n        )\n\n    def forward(self, input):\n        return self.main(input)\n    \n\nnetG = Generator(ngpu).to(device);\nnetG.apply(weights_init)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, ngpu):\n        super(Discriminator, self).__init__()\n        self.ngpu = ngpu\n        self.main = nn.Sequential(\n            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),\n            nn.LeakyReLU(0.2, inplace=True),\n            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ndf * 2),\n            nn.LeakyReLU(0.2, inplace=True),\n            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ndf * 4),\n            nn.LeakyReLU(0.2, inplace=True),\n            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ndf * 8),\n            nn.LeakyReLU(0.2, inplace=True),\n            # state size. (ndf*8) x 4 x 4\n            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),\n            nn.Sigmoid()\n        )\n\n    def forward(self, input):\n        return self.main(input)\n    \nnetD = Discriminator(ngpu).to(device);\nnetD.apply(weights_init);","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Set Loss Function"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Initialize BCELoss function\ncriterion = nn.BCELoss()\n\nfixed_noise = torch.randn(64, nz, 1, 1, device=device)\n\n# Establish convention for real and fake labels during training\nreal_label = 1\nfake_label = 0\n\n# Setup Adam optimizers for both G and D\noptimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))\noptimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))\nscheduler = ReduceLROnPlateau(optimizerD, factor=0.5, patience=2)\nscheduler = StepLR(optimizerD, gamma=0.5, step_size=8)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model training"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Training Loop\n\n# Lists to keep track of progress\n\nG_losses = []\nD_losses = []\niters = 0\n\nvalid_loss_min = np.Inf\npatience = 5\n# current number of epochs, where validation loss didn't increase\np = 0\n# whether training should be stopped\nstop = False\n\n\nprint(\"Starting Training Loop...\")\n# For each epoch\nfor epoch in range(num_epochs):\n    # For each batch in the dataloader\n    for i, data in enumerate(dataloader, 0):\n        ## Train with all-real batch\n        netD.zero_grad()\n        # Format batch\n        real_cpu = data[0].to(device)\n        b_size = real_cpu.size(0)\n        label = torch.full((b_size,), real_label, device=device)\n        # Forward pass real batch through D\n        output = netD(real_cpu).view(-1)\n        # Calculate loss on all-real batch\n        errD_real = criterion(output, label)\n        # Calculate gradients for D in backward pass\n        errD_real.backward()\n        D_x = output.mean().item()\n\n        ## Train with all-fake batch\n        # Generate batch of latent vectors\n        noise = torch.randn(b_size, nz, 1, 1, device=device)\n        # Generate fake image batch with G\n        fake = netG(noise)\n        label.fill_(fake_label)\n        # Classify all fake batch with D\n        output = netD(fake.detach()).view(-1)\n        # Calculate D's loss on the all-fake batch\n        errD_fake = criterion(output, label)\n        # Calculate the gradients for this batch\n        errD_fake.backward()\n        D_G_z1 = output.mean().item()\n        # Add the gradients from the all-real and all-fake batches\n        errD = errD_real + errD_fake\n        # Update D\n        optimizerD.step()\n\n        ############################\n        # (2) Update G network: maximize log(D(G(z)))\n        ###########################\n        netG.zero_grad()\n        label.fill_(real_label)  # fake labels are real for generator cost\n        # Since we just updated D, perform another forward pass of all-fake batch through D\n        output = netD(fake).view(-1)\n        # Calculate G's loss based on this output\n        errG = criterion(output, label)\n        # Calculate gradients for G\n        errG.backward()\n        D_G_z2 = output.mean().item()\n        # Update G\n        optimizerG.step()\n\n        # Save Losses for plotting later\n        G_losses.append(errG.item())\n        D_losses.append(errD.item())\n\n        iters += 1\n    scheduler.step(errD.item())\n    \n    print('[%d/%d][%d/%d]\\tLoss_D: %.4f\\tLoss_G: %.4f\\tD(x): %.4f\\tD(G(z)): %.4f / %.4f'\n          % (epoch, num_epochs, i, len(dataloader),\n             errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(10,5))\nplt.title(\"Generator and Discriminator Loss During Training\")\nplt.plot(G_losses,label=\"G\")\nplt.plot(D_losses,label=\"D\")\nplt.xlabel(\"iterations\")\nplt.ylabel(\"Loss\")\nplt.legend()\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Predict"},{"metadata":{"trusted":true},"cell_type":"code","source":"if not os.path.exists('../output_images'):\n    os.mkdir('../output_images')\nim_batch_size = 50\nn_images=10000\nfor i_batch in range(0, n_images, im_batch_size):\n    gen_z = torch.randn(im_batch_size, nz, 1, 1, device=device)\n    gen_images = netG(gen_z)\n    images = gen_images.to(\"cpu\").clone().detach()\n    images = images.numpy().transpose(0, 2, 3, 1)\n    for i_image in range(gen_images.size(0)):\n        save_image(gen_images[i_image, :, :, :], os.path.join('../output_images', f'image_{i_batch+i_image:05d}.png'))\n\n\nimport shutil\nshutil.make_archive('images', 'zip', '../output_images')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig = plt.figure(figsize=(25, 16))\n# display 10 images from each class\nfor i, j in enumerate(images[:32]):\n    ax = fig.add_subplot(4, 8, i + 1, xticks=[], yticks=[])\n    plt.imshow(j)","execution_count":null,"outputs":[]}],"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.5"}},"nbformat":4,"nbformat_minor":1}