Commit bc17c6a2 authored by Guillaume HEUSCH's avatar Guillaume HEUSCH
Browse files

Merge branch '4-add-gans' into 'master'

Resolve "Add GANs"

Closes #4

See merge request !4
parents 3352906d a8c468d6
Pipeline #22461 passed with stages
in 9 minutes and 14 seconds
#!/usr/bin/env python
# encoding: utf-8
import torch
import torch.nn as nn
class ConditionalGAN_generator(nn.Module):
""" Class implementating the conditional GAN generator
This network is introduced in the following publication:
Mehdi Mirza, Simon Osindero: "Conditional Generative Adversarial Nets"
Attributes
----------
ngpu : int
The number of available GPU devices
main : :py:class:`torch.nn.Sequential`
The sequential container
"""
def __init__(self, noise_dim, conditional_dim, channels=3, ngpu=1):
"""Init function
Parameters
----------
noise_dim : int
The dimension of the noise
conditional_dim : int
The dimension of the conditioning variable
channels : int
The number of channels in the image
ngpu : int
The number of available GPU devices
"""
super(ConditionalGAN_generator, self).__init__()
self.ngpu = ngpu
self.conditional_dim = conditional_dim
# output dimension
ngf = 64
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, z, y):
"""Forward function
Parameters
----------
z : :py:class: `torch.autograd.Variable`
The minibatch of noise.
y : :py:class: `torch.autograd.Variable`
The conditional one hot encoded vector for the minibatch.
Returns
-------
:py:class:`torch.Tensor`
the output of the generator (i.e. an image)
"""
generator_input = torch.cat((z, y), 1)
if isinstance(generator_input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
output = nn.parallel.data_parallel(self.main, generator_input, range(self.ngpu))
else:
output = self.main(generator_input)
return output
class ConditionalGAN_discriminator(nn.Module):
""" Class implementating the conditional GAN discriminator
Attributes
----------
conditional_dim: int
The dimension of the conditioning variable.
channels: int
The number of channels in the input image (default: 3).
ngpu : int
The number of available GPU devices
main : :py:class:`torch.nn.Sequential`
The sequential container
"""
def __init__(self, conditional_dim, channels=3, ngpu=1):
"""Init function
Parameters
----------
conditional_dim: int
The dimension of the conditioning variable.
channels: int
The number of channels in the input image (default: 3).
ngpu : int
The number of available GPU devices
"""
super(ConditionalGAN_discriminator, self).__init__()
self.conditional_dim = conditional_dim
self.ngpu = ngpu
# input dimension
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, images, y):
"""Forward function
Parameters
----------
images : :py:class: `torch.autograd.Variable`
The minibatch of input images.
y : :py:class: `torch.autograd.Variable`
The corresponding conditional feature maps.
Returns
-------
:py:class:`torch.Tensor`
the output of the discriminator
"""
input_discriminator = torch.cat((images, y), 1)
if isinstance(input_discriminator.data, torch.cuda.FloatTensor) and self.ngpu > 1:
output = nn.parallel.data_parallel(self.main, input_discriminator, range(self.ngpu))
else:
output = self.main(input_discriminator)
return output.view(-1, 1).squeeze(1)
#!/usr/bin/env python
# encoding: utf-8
import torch
import torch.nn as nn
class DCGAN_generator(nn.Module):
""" Class implementating the generator part of the Deeply Convolutional GAN
This network is introduced in the following publication:
Alec Radford, Luke Metz, Soumith Chintala: "Unsupervised Representation
Learning with Deep Convolutional Generative Adversarial Networks", ICLR 2016
and most of the code is based on:
https://github.com/pytorch/examples/tree/master/dcgan
Attributes
----------
ngpu : int
The number of available GPU devices
"""
def __init__(self, ngpu):
"""Init function
Parameters
----------
ngpu : int
The number of available GPU devices
"""
super(DCGAN_generator, self).__init__()
self.ngpu = ngpu
# just to test - will soon be args
nz = 100 # noise dimension
ngf = 64 # number of features map on the first layer
nc = 3 # number of channels
self.main = nn.Sequential(
# input is Z, going into a convolution
nn.ConvTranspose2d( nz, 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, nc, 4, 2, 1, bias=False),
nn.Tanh()
# state size. (nc) x 64 x 64
)
def forward(self, input):
"""Forward function
Parameters
----------
input : :py:class:`torch.Tensor`
Returns
-------
:py:class:`torch.Tensor`
the output of the generator (i.e. an image)
"""
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 DCGAN_discriminator(nn.Module):
""" Class implementating the discriminator part of the Deeply Convolutional GAN
This network is introduced in the following publication:
Alec Radford, Luke Metz, Soumith Chintala: "Unsupervised Representation
Learning with Deep Convolutional Generative Adversarial Networks", ICLR 2016
and most of the code is based on:
https://github.com/pytorch/examples/tree/master/dcgan
Attributes
----------
ngpu : int
The number of available GPU devices
"""
def __init__(self, ngpu):
"""Init function
Parameters
----------
ngpu : int
The number of available GPU devices
"""
super(DCGAN_discriminator, self).__init__()
self.ngpu = ngpu
# just to test - will soon be args
ndf = 64
nc = 3
self.main = nn.Sequential(
# input is (nc) x 64 x 64
nn.Conv2d(nc, 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):
"""Forward function
Parameters
----------
input : :py:class:`torch.Tensor`
Returns
-------
:py:class:`torch.Tensor`
the output of the generator (i.e. an image)
"""
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 .CNN8 import CNN8
from .CASIANet import CASIANet
from .DCGAN import DCGAN_generator
from .DCGAN import DCGAN_discriminator
from .ConditionalGAN import ConditionalGAN_generator
from .ConditionalGAN import ConditionalGAN_discriminator
from .utils import weights_init
# gets sphinx autodoc done right - don't remove it
__all__ = [_ for _ in dir() if not _.startswith('_')]
......@@ -30,3 +30,24 @@ def make_conv_layers(cfg, input_c = 3):
in_channels = v
return nn.Sequential(*layers)
def weights_init(m):
""" Initialize the weights
Initialize the weights in the different layers of
the network.
Parameters
----------
m : :py:class:`torch.nn.Conv2d`
The layer to initialize
"""
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)
......@@ -52,48 +52,6 @@ from bob.learn.pytorch.utils import get_parameter
version = pkg_resources.require('bob.learn.pytorch')[0].version
def get_parameter(args, configuration, keyword, default):
""" Get the right value for a parameter
The parameters are either defined in a separate configuration file
or given directly via command-line. Note that the command-line
has priority over the configuration file.
As a convention, parameters made with more than one word (i.e. batch size)
are provided with an underscore in the configuration file, and with an
hyphen in the command-line:
- configuration: batch_size=64
- command line: --batch-size=64
Parameters
----------
args: dictionary
The arguments as parsed from the command line.
configuration: object
The arguments given by the configuration file.
keyword: string
the keyword for the parameter to process (in the "configuration" style)
default:
The default value of the parameter
Returns
-------
arg:
The right value for the given keyword argument
"""
args_kw = '--' + keyword.replace('_', '-')
_type = type(default)
arg = _type(args[args_kw])
if hasattr(configuration, keyword):
arg = getattr(configuration, keyword)
if _type(args[args_kw]) is not default:
arg = _type(args[args_kw])
return arg
def main(user_input=None):
# Parse the command-line arguments
......
#!/usr/bin/env python
# encoding: utf-8
#!/usr/bin/env python
# encoding: utf-8
""" Train a Conditional GAN
Usage:
%(prog)s <configuration>
[--noise-dim=<int>] [--conditional-dim=<int>]
[--batch-size=<int>] [--epochs=<int>] [--sample=<int>]
[--output-dir=<path>] [--use-gpu] [--seed=<int>] [--verbose ...]
Arguments:
<configuration> A configuration file, defining the dataset and the network
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: 1e10]
-o, --output-dir=<path> Dir to save the logs, models and images [default: ./conditionalgan/]
-g, --use-gpu Use the GPU
-S, --seed=<int> The random seed [default: 3]
-v, --verbose Increase the verbosity (may appear multiple times).
Note that arguments provided directly by command-line will override the ones in the configuration file.
Example:
To run the training process
$ %(prog)s config.py
See '%(prog)s --help' for more information.
"""
import os, sys
import pkg_resources
import torch
import numpy
from docopt import docopt
import bob.core
logger = bob.core.log.setup("bob.learn.pytorch")
import bob.io.base
from bob.extension.config import load
from bob.learn.pytorch.trainers import ConditionalGANTrainer
from bob.learn.pytorch.utils import get_parameter
version = pkg_resources.require('bob.learn.pytorch')[0].version
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 conditional GAN (%s)' % version,)
# load configuration file
configuration = load([os.path.join(args['<configuration>'])])
# get various parameters, either from config file or command-line
noise_dim = get_parameter(args, configuration, 'noise_dim', 100)
conditional_dim = get_parameter(args, configuration, 'conditional_dim', 13)
batch_size = get_parameter(args, configuration, 'batch_size', 64)
epochs = get_parameter(args, configuration, 'epochs', 20)
sample = get_parameter(args, configuration, 'sample', 1e10)
seed = get_parameter(args, configuration, 'seed', 3)
output_dir = get_parameter(args, configuration, 'output_dir', 'training')
use_gpu = get_parameter(args, configuration, 'use_gpu', False)
verbosity_level = get_parameter(args, configuration, 'verbose', 0)
bob.core.log.set_verbosity_level(logger, verbosity_level)
images_dir = os.path.join(output_dir, 'samples')
log_dir = os.path.join(output_dir, 'logs')
model_dir = os.path.join(output_dir, 'models')
bob.io.base.create_directories_safe(images_dir)
bob.io.base.create_directories_safe(log_dir)
bob.io.base.create_directories_safe(images_dir)
# print parameters
logger.debug("Noise dimension = {}".format(noise_dim))
logger.debug("Conditional dimension = {}".format(conditional_dim))
logger.debug("Batch size = {}".format(batch_size))
logger.debug("Epochs = {}".format(epochs))
logger.debug("Sample = {}".format(sample))
logger.debug("Seed = {}".format(seed))
logger.debug("Output directory = {}".format(output_dir))
logger.debug("Use GPU = {}".format(use_gpu))
# 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")
# get data
if hasattr(configuration, 'dataset'):
dataloader = torch.utils.data.DataLoader(configuration.dataset, batch_size=batch_size, shuffle=True)
logger.info("There are {} training images".format(len(configuration.dataset)))
else:
logger.error("Please provide a dataset in your configuration file !")
sys.exit()
# train the model
if hasattr(configuration, 'generator') and hasattr(configuration, 'discriminator'):
trainer = ConditionalGANTrainer(configuration.generator, configuration.discriminator, [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)
else:
logger.error("Please provide both a generator and a discriminator in your configuration file !")
sys.exit()
#!/usr/bin/env python
# encoding: utf-8
""" Train a DCGAN to generate face images
Usage:
%(prog)s <configuration>
[--noise-dim=<int>]
[--batch-size=<int>] [--epochs=<int>] [--sample=<int>]
[--output-dir=<path>] [--use-gpu] [--seed=<int>] [--verbose ...]
Arguments:
<configuration> A configuration file, defining the dataset and the network
Options:
-h, --help Show this screen.
-V, --version Show version.
-n, --noise-dim=<int> the dimension of the noise [default: 100]
-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: 1e10]
-o, --output-dir=<path> Dir to save the logs, models and images [default: ./dcgan/]
-g, --use-gpu Use the GPU
-S, --seed=<int> The random seed [default: 3]
-v, --verbose Increase the verbosity (may appear multiple times).
Note that arguments provided directly by command-line will override the ones in the configuration file.
Example:
To run the training process
$ %(prog)s config.py
See '%(prog)s --help' for more information.
"""
import os, sys
import pkg_resources
import torch
import numpy
from docopt import docopt
import bob.core
logger = bob.core.log.setup("bob.learn.pytorch")
import bob.io.base
from bob.extension.config import load
from bob.learn.pytorch.trainers import DCGANTrainer
from bob.learn.pytorch.utils import get_parameter
version = pkg_resources.require('bob.learn.pytorch')[0].version
def main(user_input=None):