GAN testing

Author: spinorx

.gitignore

@@ -4,3 +4,4 @@ __pycache__
 *.lock
 xcuserdata
 *.py.cfg
+.python-version

gan/.gitignore

@@ -0,0 +1,5 @@
+data
+dataset
+__pycache__
+runs
+

gan/gan.py

@@ -0,0 +1,142 @@
+import torch
+from torch import nn, optim
+from torch.utils.data import DataLoader
+from torchvision import transforms, datasets
+from utils import Logger
+
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+def minst_data():
+  compose = transforms.Compose(
+      [transforms.ToTensor(),
+       transforms.Normalize((.5, .5, .5), (.5, .5, .5))])
+  out_dir = './dataset'
+  return datasets.MNIST(root=out_dir, train=True, transform=compose, download=True)
+
+def images_to_vectors(images):
+  return images.view(images.size(0), 784)
+
+def vectors_to_images(vectors):
+  return vectors.view(vectors.size(0), 1, 28, 28)
+
+GEN_N_FEATURES = 100
+
+def noise(size):
+  n = torch.randn(size, GEN_N_FEATURES).to(DEVICE)
+  return n
+
+def ones_target(size):
+  n = torch.ones(size, 1).to(DEVICE)
+  return n
+
+def zeros_target(size):
+  n = torch.zeros(size, 1).to(DEVICE)
+  return n
+
+class DiscriminatorNet(torch.nn.Module):
+  def __init__(self):
+    super(DiscriminatorNet, self).__init__()
+    n_features = 784
+    n_out = 1
+    self.hidden0 = nn.Sequential(
+        nn.Linear(n_features, 1024),
+        nn.LeakyReLU(0.2),
+        nn.Dropout(0.3))
+    self.hidden1 = nn.Sequential(
+        nn.Linear(1024, 512),
+        nn.LeakyReLU(0.2),
+        nn.Dropout(0.3))
+    self.hidden2 = nn.Sequential(
+        nn.Linear(512, 256),
+        nn.LeakyReLU(0.2),
+        nn.Dropout(0.3))
+    self.out = nn.Sequential(
+        nn.Linear(256, n_out),
+        nn.Sigmoid())
+
+  def forward(self, x):
+    x = self.hidden0(x)
+    x = self.hidden1(x)
+    x = self.hidden2(x)
+    x = self.out(x)
+    return x
+
+class GeneratorNet(torch.nn.Module):
+  def __init__(self):
+    super(GeneratorNet, self).__init__()
+    n_out = 784
+    self.hidden0 = nn.Sequential(
+        nn.Linear(GEN_N_FEATURES, 256),
+        nn.LeakyReLU(0.2))
+    self.hidden1 = nn.Sequential(
+        nn.Linear(256, 512),
+        nn.LeakyReLU(0.2))
+    self.hidden2 = nn.Sequential(
+        nn.Linear(512, 1024),
+        nn.LeakyReLU(0.2))
+    self.out = nn.Sequential(
+        nn.Linear(1024, n_out),
+        nn.Tanh())
+
+  def forward(self, x):
+    x = self.hidden0(x)
+    x = self.hidden1(x)
+    x = self.hidden2(x)
+    x = self.out(x)
+    return x
+
+def train_discriminator(generator, discriminator, loss, optimizer, real_data):
+  N = real_data.size(0)
+  optimizer.zero_grad()
+  # Training with real data
+  prediction_real = discriminator(real_data)
+  error_real = loss(prediction_real, ones_target(N))
+  error_real.backward()
+  # Train on Fake data
+  fake_data = generator(noise(N)).detach()
+  prediction_fake = discriminator(fake_data)
+  error_fake = loss(prediction_fake, zeros_target(N))
+  error_fake.backward()
+
+  optimizer.step()
+  return error_real + error_fake, prediction_real, prediction_fake
+
+def train_generator(generator, discriminator, loss, optimizer, N):
+  optimizer.zero_grad()
+  fake_data = generator(noise(N))
+  prediction = discriminator(fake_data)
+  error = loss(prediction, ones_target(N))
+  error.backward()
+  optimizer.step()
+  return error
+
+if __name__ == '__main__':
+  data = minst_data()
+  data_loader = DataLoader(data, batch_size=100, shuffle=True)
+  num_batches = len(data_loader)
+  discriminator = DiscriminatorNet().to(DEVICE)
+  generator = GeneratorNet().to(DEVICE)
+  d_optimizer = optim.Adam(discriminator.parameters(), lr=0.0002)
+  g_optimizer = optim.Adam(generator.parameters(), lr=0.0002)
+  loss = nn.BCELoss()
+
+  num_test_samples = 16
+  test_noise = noise(num_test_samples)
+
+  logger = Logger(model_name='VGAN', data_name='MNIST')
+  num_epochs = 200
+
+  for epoch in range(num_epochs):
+    for n_batch, (real_batch,_) in enumerate(data_loader):
+      N = real_batch.size(0)
+      # Train Discriminator
+      real_data = images_to_vectors(real_batch).to(DEVICE)
+      d_error, d_pred_real, d_pred_fake = train_discriminator(
+          generator, discriminator, loss, d_optimizer, real_data)
+      g_error = train_generator(generator, discriminator, loss, g_optimizer, N)
+      logger.log(d_error, g_error, epoch, n_batch, num_batches)
+      if n_batch % 100 == 0:
+        test_images = vectors_to_images(generator(test_noise)).cpu()
+        logger.log_images(test_images.data, num_test_samples, epoch, n_batch, num_batches)
+        logger.display_status(epoch, num_epochs, n_batch, num_batches, d_error, g_error, d_pred_real, d_pred_fake)
+

gan/utils.py

@@ -0,0 +1,136 @@
+import os
+import numpy as np
+import errno
+import torchvision.utils as vutils
+from tensorboardX import SummaryWriter
+from IPython import display
+from matplotlib import pyplot as plt
+import torch
+
+'''
+  TensorBoard Data will be stored in './runs' path
+'''
+
+class Logger:
+  def __init__(self, model_name, data_name):
+    self.model_name = model_name
+    self.data_name = data_name
+
+    self.comment = '{}_{}'.format(model_name, data_name)
+    self.data_subdir = '{}/{}'.format(model_name, data_name)
+
+    # TensorBoard
+    self.writer = SummaryWriter(comment=self.comment)
+
+  def log(self, d_error, g_error, epoch, n_batch, num_batches):
+
+    # var_class = torch.autograd.variable.Variable
+    if isinstance(d_error, torch.autograd.Variable):
+      d_error = d_error.data.cpu().numpy()
+    if isinstance(g_error, torch.autograd.Variable):
+      g_error = g_error.data.cpu().numpy()
+
+    step = Logger._step(epoch, n_batch, num_batches)
+    self.writer.add_scalar(
+        '{}/D_error'.format(self.comment), d_error, step)
+    self.writer.add_scalar(
+        '{}/G_error'.format(self.comment), g_error, step)
+
+  def log_images(self, images, num_images, epoch, n_batch, num_batches, format='NCHW', normalize=True):
+    '''
+    input images are expected in format (NCHW)
+    '''
+    if type(images) == np.ndarray:
+      images = torch.from_numpy(images)
+
+    if format=='NHWC':
+      images = images.transpose(1,3)
+
+
+    step = Logger._step(epoch, n_batch, num_batches)
+    img_name = '{}/images{}'.format(self.comment, '')
+
+    # Make horizontal grid from image tensor
+    horizontal_grid = vutils.make_grid(
+        images, normalize=normalize, scale_each=True)
+    # Make vertical grid from image tensor
+    nrows = int(np.sqrt(num_images))
+    grid = vutils.make_grid(
+        images, nrow=nrows, normalize=True, scale_each=True)
+
+    # Add horizontal images to tensorboard
+    self.writer.add_image(img_name, horizontal_grid, step)
+
+    # Save plots
+    self.save_torch_images(horizontal_grid, grid, epoch, n_batch)
+
+  def save_torch_images(self, horizontal_grid, grid, epoch, n_batch, plot_horizontal=True):
+    out_dir = './data/images/{}'.format(self.data_subdir)
+    Logger._make_dir(out_dir)
+
+    # Plot and save horizontal
+    fig = plt.figure(figsize=(16, 16))
+    plt.imshow(np.moveaxis(horizontal_grid.numpy(), 0, -1))
+    plt.axis('off')
+    if plot_horizontal:
+      display.display(plt.gcf())
+    self._save_images(fig, epoch, n_batch, 'hori')
+    plt.close()
+
+    # Save squared
+    fig = plt.figure()
+    plt.imshow(np.moveaxis(grid.numpy(), 0, -1))
+    plt.axis('off')
+    self._save_images(fig, epoch, n_batch)
+    plt.close()
+
+  def _save_images(self, fig, epoch, n_batch, comment=''):
+    out_dir = './data/images/{}'.format(self.data_subdir)
+    Logger._make_dir(out_dir)
+    fig.savefig('{}/{}_epoch_{}_batch_{}.png'.format(
+        out_dir, comment, epoch, n_batch))
+
+  def display_status(self, epoch, num_epochs, n_batch, num_batches, d_error, g_error, d_pred_real, d_pred_fake):
+
+    # var_class = torch.autograd.variable.Variable
+    if isinstance(d_error, torch.autograd.Variable):
+      d_error = d_error.data.cpu().numpy()
+    if isinstance(g_error, torch.autograd.Variable):
+      g_error = g_error.data.cpu().numpy()
+    if isinstance(d_pred_real, torch.autograd.Variable):
+      d_pred_real = d_pred_real.data
+    if isinstance(d_pred_fake, torch.autograd.Variable):
+      d_pred_fake = d_pred_fake.data
+
+
+    print('Epoch: [{}/{}], Batch Num: [{}/{}]'.format(
+        epoch,num_epochs, n_batch, num_batches))
+    print('Discriminator Loss: {:.4f}, Generator Loss: {:.4f}'.format(d_error, g_error))
+    print('D(x): {:.4f}, D(G(z)): {:.4f}'.format(d_pred_real.mean(), d_pred_fake.mean()))
+
+  def save_models(self, generator, discriminator, epoch):
+    out_dir = './data/models/{}'.format(self.data_subdir)
+    Logger._make_dir(out_dir)
+    torch.save(
+        generator.state_dict(),
+        '{}/G_epoch_{}'.format(out_dir, epoch))
+    torch.save(
+        discriminator.state_dict(),
+        '{}/D_epoch_{}'.format(out_dir, epoch))
+
+  def close(self):
+    self.writer.close()
+
+  # Private Functionality
+
+  @staticmethod
+  def _step(epoch, n_batch, num_batches):
+    return epoch * num_batches + n_batch
+
+  @staticmethod
+  def _make_dir(directory):
+    try:
+      os.makedirs(directory)
+    except OSError as e:
+      if e.errno != errno.EEXIST:
+        raise