Commit c1b4c3a7 authored by Guillaume HEUSCH's avatar Guillaume HEUSCH

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

parent 5ec40bd3
#!/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)
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
......
#!/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)
#!/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))
from .DCGANTrainer import DCGANTrainer
from .ConditionalGANTrainer import ConditionalGANTrainer
# gets sphinx autodoc done right - don't remove it
__all__ = [_ for _ in dir() if not _.startswith('_')]
......@@ -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',
],
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment