<jupyter_start><jupyter_text>Generating terrain using a Deep Convolutional GAN<jupyter_code>from __future__ import print_function

import argparse
import os
import random

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data

import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

import requests

from IPython.display import HTML<jupyter_output><empty_output><jupyter_text>Set random function seeds for reproducability.<jupyter_code>seed = 999

#manualSeed = random.randint(1, 10000)
random.seed(seed)
torch.manual_seed(seed)

print(seed)<jupyter_output>999<jupyter_text>Define data directory and the number of dataloader workers.<jupyter_code>data_directory = "data/Terrain/processed"
n_workers = 2<jupyter_output><empty_output><jupyter_text>Settings and hyperparameters<jupyter_code>batch_size = 8
image_size = 64

n_channels = 1
n_latent   = 100

n_gen_features = 64
n_dis_features = 64

n_epochs    = 100
learning_r  = 0.0002
beta1       = 0.5

n_gpus = 2<jupyter_output><empty_output><jupyter_text>Generator<jupyter_code>class Generator(nn.Module):
    def __init__(self, n_gpus):
        super(Generator, self).__init__()
        self.n_gpus = n_gpus
        self.main = nn.Sequential(
            nn.ConvTranspose2d(n_latent, n_gen_features * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(n_gen_features * 8),
            nn.ReLU(True),

            nn.ConvTranspose2d(n_gen_features * 8, n_gen_features * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(n_gen_features * 4),
            nn.ReLU(True),

            nn.ConvTranspose2d(n_gen_features * 4, n_gen_features * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(n_gen_features * 2),
            nn.ReLU(True),

            nn.ConvTranspose2d(n_gen_features * 2, n_gen_features, 4, 2, 1, bias=False),
            nn.BatchNorm2d(n_gen_features),
            nn.ReLU(True),

            nn.ConvTranspose2d(n_gen_features, n_channels, 4, 2, 1, bias=False),
            nn.Tanh()
        )

    def forward(self, input):
        return self.main(input)<jupyter_output><empty_output><jupyter_text>Discriminator<jupyter_code>class Discriminator(nn.Module):
    def __init__(self, n_gpus):
        super(Discriminator, self).__init__()
        self.n_gpus = n_gpus
        self.main = nn.Sequential(
            nn.Conv2d(n_channels, n_dis_features, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(n_dis_features, n_dis_features * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(n_dis_features * 2),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(n_dis_features * 2, n_dis_features * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(n_dis_features * 4),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(n_dis_features * 4, n_dis_features * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(n_dis_features * 8),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(n_dis_features * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        return self.main(input)<jupyter_output><empty_output><jupyter_text>Loading the dataset<jupyter_code>dataset_transforms = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize(image_size),
    transforms.CenterCrop(image_size),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,)),
])

dataset = dset.ImageFolder(root=data_directory, transform=dataset_transforms)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
                                         shuffle=True, num_workers=n_workers)
run_cuda = torch.cuda.is_available() and n_gpus > 0 
device   = torch.device("cuda:0" if run_cuda else "cpu")<jupyter_output><empty_output><jupyter_text>Dataset preview<jupyter_code>real_batch = next(iter(dataloader))
plt.figure(figsize=(8,8))
plt.axis("off")
plt.title("Training Images")
plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))
plt.show()<jupyter_output><empty_output><jupyter_text>Function for weight initialization<jupyter_code>def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)

def save_img(image, name):
    array = array.detach()
    array = array.reshape(28,28)
    
    plt.imshow(immage, cmap='binary')
    plt.xticks([])
    plt.yticks([])
    plt.savefig('output/' + name)<jupyter_output><empty_output><jupyter_text>Generator and discriminator definition and initialization<jupyter_code>generator = Generator(n_gpus).to(device)
discriminator = Discriminator(n_gpus).to(device)

if device.type is 'cuda' and n_gpus > 1:
    generator = nn.DataParallel(n_gpus, list(range(n_gpus)))
    discriminator = nn.DataParallel(n_gpus, list(range(n_gpus)))

generator.apply(weights_init)
discriminator.apply(weights_init)<jupyter_output><empty_output><jupyter_text>Loss function and optimizer initialization<jupyter_code>criterion = nn.BCELoss()

fixed_noise = torch.randn(batch_size, n_latent, 1, 1, device=device)

real_label = 1
fake_label = 0

d_optimizer = optim.Adam(discriminator.parameters(), lr=learning_r, betas=(beta1, 0.999))
g_optimizer = optim.Adam(generator.parameters(), lr=learning_r, betas=(beta1, 0.999))<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>img_progress = []
gen_loss, dis_loss = [], []
iters = 0
# Start training
for epoch in range(n_epochs):
    # requests.get('http://localhost:5000/gan/vars/epochs/{}%2F{}'.format(epoch, n_epochs))

    for i, data in enumerate(dataloader, 0):
        # Train discriminator
        discriminator.zero_grad()

        real_cpu = data[0].to(device)
        b_size = real_cpu.size(0)

        # Train discriminator with real batch
        label = torch.full((b_size,), real_label, device=device)

        output = discriminator(real_cpu).view(-1)

        d_error_r = criterion(output, label)
        d_error_r.backward(retain_graph=True)

        # Train discriminator with fake batch
        noise = torch.randn(b_size, n_latent, 1, 1, device=device)

        fake = generator(noise)
        label.fill_(fake_label)

        output = discriminator(fake.detach())

        d_error_f = criterion(output, label)
        d_error_f.backward(retain_graph=True)

        # Step discriminator
        d_error = d_error_r + d_error_f
        d_optimizer.step()

        # Train generator
        generator.zero_grad()
        label.fill_(real_label)

        output = discriminator(fake).view(-1)

        g_error = criterion(output, label)
        g_error.backward(retain_graph=True)

        g_optimizer.step()

        if i % 50 is 0:
            # requests.get('http://localhost:5000/gan/vars/iterations/{}%2F{}'.format(i, len(dataloader)))
            # requests.get('http://localhost:5000/gan/vars/gen_loss/{}'.format(g_error.item()))
            # requests.get('http://localhost:5000/gan/vars/dis_loss/{}'.format(d_error.item()))

            print('[%d/%d][%d/%d]\t d_loss: %.4f\t g_loss: %.4f\t cuda: %s' %
                (epoch + 1, n_epochs, i, len(dataloader), d_error.item(), g_error.item(), run_cuda))
        
        gen_loss.append(g_error.item())
        dis_loss.append(d_error.item())

        # Check how the generator is doing by saving G's output on fixed_noise
        if iters % 1000 is 0 or (epoch is n_epochs - 1 and i == len(dataloader) - 1):
            torch.save(generator, 'saved_models/generator_{}.pt'.format(iters))
            torch.save(generator, 'saved_models/generator_latest.pt')

            torch.save(discriminator, 'saved_models/discriminator_{}.pt'.format(iters))
            torch.save(discriminator, 'saved_models/discriminator_latest.pt')

            with torch.no_grad():
                fake = generator(fixed_noise).detach().cpu()
            img_progress.append(vutils.make_grid(fake, padding=2, normalize=True))
            plt.imsave('output/terrain_{}.png'.format(iters), img_progress[-1][0])
            plt.imsave('output/terrain_gray_{}.png'.format(iters), img_progress[-1][0], cmap='Greys')

        iters += 1
plt.figure(figsize=(10,5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(gen_loss,label="G")
plt.plot(dis_loss,label="D")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.show()
fig = plt.figure(figsize=(8,8))
plt.axis("off")
images = [[plt.imshow(np.transpose(i, (1, 2, 0)), animated=True)] for i in img_progress]
anim = animation.ArtistAnimation(fig, images, interval=1000, repeat_delay=1000, blit=True)

HTML(anim.to_jshtml())<jupyter_output><empty_output><jupyter_text>Generate!<jupyter_code>latent = torch.randn(1, n_latent, 1, 1, device=device)

with torch.no_grad():
    fake = generator(latent).cpu()
plt.imshow(vutils.make_grid(fake, padding=2, normalize=True)[0])<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup