[all] added the code for Conditional GAN (work in progress)

bob/learn/pytorch/architectures/ConditionalGAN.py

+#!/usr/bin/env python
+# encoding: utf-8
+
+
+import torch
+import torch.nn as nn
+
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        m.weight.data.normal_(0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        m.weight.data.normal_(1.0, 0.02)
+        m.bias.data.fill_(0)
+
+
+class ConditionalGAN_generator(nn.Module):
+  def __init__(self, noise_dim, conditional_dim, channels=3, ngpu=1):
+    super(ConditionalGAN_generator, self).__init__()
+    self.ngpu = ngpu
+        
+        
+    # just to test - will soon be args
+    
+    ngf = 64
+    print type(noise_dim)
+
+    self.main = nn.Sequential(
+      # input is Z, going into a convolution
+      nn.ConvTranspose2d((noise_dim + conditional_dim), ngf * 8, 4, 1, 0, bias=False),
+      nn.BatchNorm2d(ngf * 8),
+      nn.ReLU(True),
+      # state size. (ngf*8) x 4 x 4
+      nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
+      nn.BatchNorm2d(ngf * 4),
+      nn.ReLU(True),
+      # state size. (ngf*4) x 8 x 8
+      nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
+      nn.BatchNorm2d(ngf * 2),
+      nn.ReLU(True),
+      # state size. (ngf*2) x 16 x 16
+      nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
+      nn.BatchNorm2d(ngf),
+      nn.ReLU(True),
+      # state size. (ngf) x 32 x 32
+      nn.ConvTranspose2d(ngf, channels, 4, 2, 1, bias=False),
+      nn.Tanh()
+      # state size. (nc) x 64 x 64
+    )
+
+  def forward(self, input):
+    if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
+      output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
+    else:
+      output = self.main(input)
+    return output
+
+class ConditionalGAN_discriminator(nn.Module):
+  def __init__(self, conditional_dim, channels=3, ngpu=1):
+    super(ConditionalGAN_discriminator, self).__init__()
+    self.ngpu = ngpu
+        
+        
+    # just to test - will soon be args
+    ndf = 64
+       
+    self.main = nn.Sequential(
+      # input is (nc) x 64 x 64
+      nn.Conv2d((channels + conditional_dim), ndf, 4, 2, 1, bias=False),
+      nn.LeakyReLU(0.2, inplace=True),
+      # state size. (ndf) x 32 x 32
+      nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
+      nn.BatchNorm2d(ndf * 2),
+      nn.LeakyReLU(0.2, inplace=True),
+      # state size. (ndf*2) x 16 x 16
+      nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
+      nn.BatchNorm2d(ndf * 4),
+      nn.LeakyReLU(0.2, inplace=True),
+      # state size. (ndf*4) x 8 x 8
+      nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
+      nn.BatchNorm2d(ndf * 8),
+      nn.LeakyReLU(0.2, inplace=True),
+      # state size. (ndf*8) x 4 x 4
+      nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
+      nn.Sigmoid()
+    )
+
+  def forward(self, input):
+    if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
+      output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
+    else:
+      output = self.main(input)
+
+    return output.view(-1, 1).squeeze(1)

bob/learn/pytorch/architectures/__init__.py

 from .DCGAN import DCGAN_generator
 from .DCGAN import DCGAN_discriminator
+
+from .ConditionalGAN import ConditionalGAN_generator
+from .ConditionalGAN import ConditionalGAN_discriminator
+
 from .DCGAN import weights_init
 
 # gets sphinx autodoc done right - don't remove it

bob/learn/pytorch/scripts/train_conditionalgan_multipie.py

+#!/usr/bin/env python
+# encoding: utf-8
+
+
+""" Train a Conditional GAN 
+
+Usage:
+  %(prog)s [--noise-dim=<int>] [--conditional-dim=<int>] 
+           [--batch-size=<int>] [--epochs=<int>] [--sample=<int>]
+           [--output-dir=<path>] [--use-gpu] [--seed=<int>] [--verbose ...]
+
+Options:
+  -h, --help                    Show this screen.
+  -V, --version                 Show version.
+  -n, --noise-dim=<int>         The dimension of the noise [default: 100]
+  -c, --conditional-dim=<int>   The dimension of the conditional variable [default: 13]
+  -b, --batch-size=<int>        The size of your mini-batch [default: 64]
+  -e, --epochs=<int>            The number of training epochs [default: 100] 
+  -s, --sample=<int>            Save generated images at every 'sample' batch iteration [default: 100000000000] 
+  -o, --output-dir=<path>       Dir to save the logs, models and images [default: ./drgan-light-mpie-casia/] 
+  -g, --use-gpu                 Use the GPU 
+  -S, --seed=<int>              The random seed [default: 3] 
+  -v, --verbose                 Increase the verbosity (may appear multiple times).
+
+Example:
+
+  To run the training process 
+
+    $ %(prog)s --batch-size 64 --epochs 25 --output-dir drgan 
+
+See '%(prog)s --help' for more information.
+
+"""
+
+import os, sys
+import pkg_resources
+
+import bob.core
+logger = bob.core.log.setup("bob.learn.pytorch")
+
+from docopt import docopt
+
+version = pkg_resources.require('bob.learn.pytorch')[0].version
+
+import numpy
+import bob.io.base
+
+# torch
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision.transforms as transforms
+import torchvision.utils as vutils
+from torch.autograd import Variable
+
+# data and architecture from the package
+from bob.learn.pytorch.datasets import MultiPIEDataset
+from bob.learn.pytorch.datasets import RollChannels
+from bob.learn.pytorch.datasets import ToTensor
+from bob.learn.pytorch.datasets import Normalize
+
+from bob.learn.pytorch.architectures import ConditionalGAN_generator
+from bob.learn.pytorch.architectures import ConditionalGAN_discriminator
+from bob.learn.pytorch.architectures import weights_init
+
+from bob.learn.pytorch.trainers import ConditionalGANTrainer
+
+
+def main(user_input=None):
+  
+  # Parse the command-line arguments
+  if user_input is not None:
+      arguments = user_input
+  else:
+      arguments = sys.argv[1:]
+
+  prog = os.path.basename(sys.argv[0])
+  completions = dict(prog=prog, version=version,)
+  args = docopt(__doc__ % completions,argv=arguments,version='Train DR-GAN (%s)' % version,)
+
+  # verbosity
+  verbosity_level = args['--verbose']
+  bob.core.log.set_verbosity_level(logger, verbosity_level)
+
+  # get the arguments
+  noise_dim = int(args['--noise-dim'])
+  conditional_dim = int(args['--conditional-dim'])
+  batch_size = int(args['--batch-size'])
+  epochs = int(args['--epochs'])
+  sample = int(args['--sample'])
+  output_dir = str(args['--output-dir'])
+  seed = int(args['--seed'])
+  use_gpu = bool(args['--use-gpu'])
+
+  images_dir = os.path.join(output_dir, 'samples')
+  log_dir = os.path.join(output_dir, 'logs')
+  model_dir = os.path.join(output_dir, 'models')
+
+  # process on the arguments / options
+  torch.manual_seed(seed)
+  if use_gpu:
+    torch.cuda.manual_seed_all(seed)
+  if torch.cuda.is_available() and not use_gpu:
+    logger.warn("You have a CUDA device, so you should probably run with --use-gpu")
+  bob.io.base.create_directories_safe(images_dir)
+  bob.io.base.create_directories_safe(log_dir)
+  bob.io.base.create_directories_safe(images_dir)
+
+  # ============
+  # === DATA ===
+  # ============
+  
+  # WARNING with the transforms ... act on labels too, at some point, I may have to write my own
+  # Also, in 'ToTensor', there is a reshape performed from: HxWxC to CxHxW
+  face_dataset = MultiPIEDataset(root_dir='/idiap/temp/heusch/data/multipie-cropped-64x64', 
+                                 frontal_only=False, 
+                                 #transform=None
+                                 transform=transforms.Compose([
+                                   RollChannels(), # bob to skimage:  
+                                   ToTensor(),
+                                   Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+                                   ])
+                                 )
+
+  dataloader = torch.utils.data.DataLoader(face_dataset, batch_size=batch_size, shuffle=True)
+  logger.info("There are {} training images".format(len(face_dataset))) 
+  
+  # ===============
+  # === NETWORK ===
+  # ===============
+  ngpu = 1 # usually we don't have more than one GPU
+  
+  generator = ConditionalGAN_generator(noise_dim, conditional_dim)
+  generator.apply(weights_init)
+  logger.info("Generator architecture: {}".format(generator))
+
+  discrminator = ConditionalGAN_discriminator(conditional_dim)
+  discrminator.apply(weights_init)
+  logger.info("Discriminator architecture: {}".format(discrminator))
+
+
+  # ===============
+  # === TRAINER ===
+  # ===============
+  trainer = ConditionalGANTrainer(generator, discrminator, [3, 64, 64], batch_size=batch_size, noise_dim=noise_dim, conditional_dim=conditional_dim, use_gpu=use_gpu, verbosity_level=verbosity_level)
+  trainer.train(dataloader, n_epochs=epochs, output_dir=output_dir)

bob/learn/pytorch/trainers/ConditionalGANTrainer.py

+#!/usr/bin/env python
+# encoding: utf-8
+
+import numpy
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.autograd import Variable
+import torchvision.utils as vutils
+
+
+import bob.core
+logger = bob.core.log.setup("bob.learn.pytorch")
+
+import time
+
+
+class ConditionalGANTrainer(object):
+  """
+  Class to train a Conditional GAN
+
+  **Parameters**
+
+  generator: pytorch nn.Module
+    The generator network
+
+  discriminator: pytorch nn.Module
+    The discriminator network
+
+  image_size: list
+    The size of the images in this format: [channels,height, width]
+
+  batch_size: int
+    The size of your minibatch
+
+  noise_dim: int
+    The dimension of the noise (input to the generator)
+
+  conditional_dim: int
+    The dimension of the conditioning variable
+
+  use_gpu: boolean
+    If you would like to use the gpu
+
+  verbosity_level: int
+    The level of verbosity output to stdout
+  """
+  def __init__(self, netG, netD, image_size, batch_size=64, noise_dim=100, conditional_dim=13, use_gpu=False, verbosity_level=2):
+   
+    self.netG = netG
+    self.netD = netD
+    self.image_size = image_size
+    self.batch_size = batch_size
+    self.noise_dim = noise_dim 
+    self.conditional_dim = conditional_dim 
+    self.use_gpu = use_gpu
+
+    self.input = torch.FloatTensor(batch_size, (image_size[0] + conditional_dim), image_size[1], image_size[2])
+    self.conditional_noise = torch.FloatTensor(batch_size, noise_dim + conditional_dim, 1, 1)
+    self.label = torch.FloatTensor(batch_size)
+
+    # to generate samples
+    self.fixed_noise = torch.FloatTensor(12, noise_dim, 1, 1).normal_(0, 1)
+    self.fixed_noise = Variable(self.fixed_noise)
+    
+    self.criterion = nn.BCELoss()
+
+    if self.use_gpu:
+      self.netD.cuda()
+      self.netG.cuda()
+      self.criterion.cuda()
+      self.input, self.label = self.input.cuda(), self.label.cuda()
+      self.noise, self.fixed_noise = self.noise.cuda(), self.fixed_noise.cuda()
+
+    bob.core.log.set_verbosity_level(logger, verbosity_level)
+
+
+  def train(self, dataloader, n_epochs=10, learning_rate=0.0002, beta1=0.5, output_dir='out'):
+    """
+    Function that performs the training.
+
+    **Parameters**
+
+    dataloader: pytorch DataLoader
+      The dataloader for your data
+
+    n_epochs: int
+      The number of epochs you would like to train for
+
+    learning_rate: float
+      The learning rate for Adam optimizer
+
+    beta1: float
+      The beta1 for Adam optimizer
+
+    output_dir: path
+      The directory where you would like to output images and models
+    """
+    real_label = 1
+    fake_label = 0
+    
+    # setup optimizer
+    optimizerD = optim.Adam(self.netD.parameters(), lr=learning_rate, betas=(beta1, 0.999))
+    optimizerG = optim.Adam(self.netG.parameters(), lr=learning_rate, betas=(beta1, 0.999))
+
+    for epoch in range(n_epochs):
+      for i, data in enumerate(dataloader, 0):
+     
+        start = time.time()
+
+        # ===========================================================
+        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
+        # ===========================================================
+      
+        # train with real
+        self.netD.zero_grad()
+        real_images = data['image']
+        poses = data['pose']
+        image_size = real_images[1].size()
+
+        if self.use_gpu:
+          real_images = real_images.cuda()
+       
+        # build the additional feature maps corresponding to the conditioning variable 
+        temp = real_images.numpy()
+        cm = numpy.zeros((self.batch_size, self.conditional_dim, temp.shape[2], temp.shape[3]))
+        for k in range(self.batch_size):
+          cm[:, int(poses[k]), :, :] = 1
+        conditional_maps = torch.FloatTensor(cm)
+        
+        # append the conditional feature maps to the original images
+        input_discriminator = torch.FloatTensor(self.batch_size, (image_size[0] + self.conditional_dim), image_size[1], image_size[2])
+        for k in range(self.batch_size):
+          input_discriminator[k] = torch.cat((real_images[k], conditional_maps[k]), 0) 
+
+
+        self.input.resize_as_(input_discriminator).copy_(input_discriminator)
+        inputv = Variable(self.input)
+        labelv = Variable(self.label)
+        output = self.netD(inputv)
+        errD_real = self.criterion(output, labelv)
+        errD_real.backward()
+        D_x = output.data.mean()
+
+        # train with fake
+        noise = torch.FloatTensor(self.batch_size, self.noise_dim)
+        noise.normal_(0, 1)
+        
+        # generate the one hot pose encoding
+        oh = numpy.zeros((self.batch_size, self.conditional_dim))
+        for k in range(self.batch_size):
+          oh[k, int(poses[k])] = 1
+        one_hot = torch.FloatTensor(oh)
+        
+        # concatenate that with the noise
+        input_generator = torch.FloatTensor(self.batch_size, (self.noise_dim + self.conditional_dim))
+        for k in range(self.batch_size):
+          input_generator[k] = torch.cat((noise[k], one_hot[k]), 0) 
+       
+        self.conditional_noise.resize_(self.batch_size, (self.noise_dim + self.conditional_dim), 1, 1).copy_(input_generator)
+        noisev = Variable(self.conditional_noise)
+        fake = self.netG(noisev)
+        
+        # build conditional fakes
+        fake_images = fake.data 
+        input_discriminator_fake = torch.FloatTensor(self.batch_size, (image_size[0] + self.conditional_dim), image_size[1], image_size[2])
+        for k in range(self.batch_size):
+          input_discriminator_fake[k] = torch.cat((fake_images[k], conditional_maps[k]), 0) 
+
+        fake_input_v = Variable(input_discriminator_fake)
+        labelv = Variable(self.label.fill_(fake_label))
+        output = self.netD(fake_input_v)
+        errD_fake = self.criterion(output, labelv)
+        errD_fake.backward()
+        D_G_z1 = output.data.mean()
+        errD = errD_real + errD_fake
+        optimizerD.step()
+
+        # =========================================
+        # (2) Update G network: maximize log(D(G(z)))
+        # =========================================
+        self.netG.zero_grad()
+        labelv = Variable(self.label.fill_(real_label))  # fake labels are real for generator cost
+        output = self.netD(fake_input_v)
+        errG = self.criterion(output, labelv)
+        errG.backward()
+        D_G_z2 = output.data.mean()
+        optimizerG.step()
+
+        end = time.time()
+        logger.info("[{}/{}][{}/{}] => Loss D = {} -- Loss G = {} (time spent: {})".format(epoch, n_epochs, i, len(dataloader), errD.data[0], errG.data[0], (end-start)))
+      
+      # save generated images at every epoch
+      #input_generator_examples = 
+      #poses = range(self.conditional_dim)
+      #for pose in pose:
+      #  oh = numpy.zeros(self.conditional_dim)
+      #  oh[pose] = 1
+      #  one_hot = torch.FloatTensor(oh)
+
+
+      #fake = self.netG(self.fixed_noise)
+
+      #vutils.save_image(fake.data, '%s/fake_samples_epoch_%03d.png' % (output_dir, epoch), normalize=True)
+
+      ## do checkpointing
+      #torch.save(self.netG.state_dict(), '%s/netG_epoch_%d.pth' % (output_dir, epoch))
+      #torch.save(self.netD.state_dict(), '%s/netD_epoch_%d.pth' % (output_dir, epoch))

bob/learn/pytorch/trainers/__init__.py

+from .DCGANTrainer import DCGANTrainer
+from .ConditionalGANTrainer import ConditionalGANTrainer
+
 # gets sphinx autodoc done right - don't remove it
 __all__ = [_ for _ in dir() if not _.startswith('_')]
 

setup.py

@@ -75,6 +75,7 @@ setup(
       # scripts should be declared using this entry:
       'console_scripts': [
         'train_dcgan_multipie.py = bob.learn.pytorch.scripts.train_dcgan_multipie:main', 
+        'train_conditionalgan_multipie.py = bob.learn.pytorch.scripts.train_conditionalgan_multipie:main', 
       ],