Skip to content
Snippets Groups Projects
Commit 0e88af7a authored by Guillaume HEUSCH's avatar Guillaume HEUSCH
Browse files

[trainers, scripts] fixed stuff in the script to train DCGAN (transforms...

[trainers, scripts] fixed stuff in the script to train DCGAN (transforms mostly), reflect changes in the DCGAN trainer
parent 3b7be9b0
No related branches found
No related tags found
No related merge requests found
...@@ -5,19 +5,18 @@ ...@@ -5,19 +5,18 @@
""" Train a DR-GAN """ Train a DR-GAN
Usage: Usage:
%(prog)s [--latent-dim=<int>] [--noise-dim=<int>] %(prog)s [--noise-dim=<int>]
[--batch-size=<int>] [--epochs=<int>] [--sample=<int>] [--batch-size=<int>] [--epochs=<int>] [--sample=<int>]
[--output-dir=<path>] [--use-gpu] [--seed=<int>] [--verbose ...] [--output-dir=<path>] [--use-gpu] [--seed=<int>] [--verbose ...]
Options: Options:
-h, --help Show this screen. -h, --help Show this screen.
-V, --version Show version. -V, --version Show version.
-l, --latent-dim=<int> the dimension of the encoded ID [default: 320]
-n, --noise-dim=<int> the dimension of the noise [default: 100] -n, --noise-dim=<int> the dimension of the noise [default: 100]
-b, --batch-size=<int> The size of your mini-batch [default: 64] -b, --batch-size=<int> The size of your mini-batch [default: 64]
-e, --epochs=<int> The number of training epochs [default: 100] -e, --epochs=<int> The number of training epochs [default: 100]
-s, --sample=<int> Save generated images at every 'sample' batch iteration [default: 100000000000] -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/] -o, --output-dir=<path> Dir to save the logs, models and images [default: ./dcgan-multipie/]
-g, --use-gpu Use the GPU -g, --use-gpu Use the GPU
-S, --seed=<int> The random seed [default: 3] -S, --seed=<int> The random seed [default: 3]
-v, --verbose Increase the verbosity (may appear multiple times). -v, --verbose Increase the verbosity (may appear multiple times).
...@@ -26,7 +25,7 @@ Example: ...@@ -26,7 +25,7 @@ Example:
To run the training process To run the training process
$ %(prog)s --batch-size 64 --epochs 25 --output-dir drgan $ %(prog)s --batch-size 64 --epochs 25 --output-dir dcgan
See '%(prog)s --help' for more information. See '%(prog)s --help' for more information.
...@@ -53,9 +52,11 @@ import torchvision.transforms as transforms ...@@ -53,9 +52,11 @@ import torchvision.transforms as transforms
import torchvision.utils as vutils import torchvision.utils as vutils
from torch.autograd import Variable from torch.autograd import Variable
# data and architecture from the package # data, architecture and trainer from the package
from bob.learn.pytorch.datasets import MultiPIEDataset from bob.learn.pytorch.datasets import MultiPIEDataset
from bob.learn.pytorch.datasets import RollChannels 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 DCGAN_generator from bob.learn.pytorch.architectures import DCGAN_generator
from bob.learn.pytorch.architectures import DCGAN_discriminator from bob.learn.pytorch.architectures import DCGAN_discriminator
...@@ -74,7 +75,7 @@ def main(user_input=None): ...@@ -74,7 +75,7 @@ def main(user_input=None):
prog = os.path.basename(sys.argv[0]) prog = os.path.basename(sys.argv[0])
completions = dict(prog=prog, version=version,) completions = dict(prog=prog, version=version,)
args = docopt(__doc__ % completions,argv=arguments,version='Train DR-GAN (%s)' % version,) args = docopt(__doc__ % completions,argv=arguments,version='Train DCGAN (%s)' % version,)
# verbosity # verbosity
verbosity_level = args['--verbose'] verbosity_level = args['--verbose']
...@@ -89,32 +90,24 @@ def main(user_input=None): ...@@ -89,32 +90,24 @@ def main(user_input=None):
seed = int(args['--seed']) seed = int(args['--seed'])
use_gpu = bool(args['--use-gpu']) 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 # process on the arguments / options
torch.manual_seed(seed) torch.manual_seed(seed)
if use_gpu: if use_gpu:
torch.cuda.manual_seed_all(seed) torch.cuda.manual_seed_all(seed)
if torch.cuda.is_available() and not use_gpu: if torch.cuda.is_available() and not use_gpu:
logger.warn("You have a CUDA device, so you should probably run with --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(output_dir)
bob.io.base.create_directories_safe(images_dir)
# ============ # ============
# === DATA === # === 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', face_dataset = MultiPIEDataset(root_dir='/idiap/temp/heusch/data/multipie-cropped-64x64',
frontal_only=True, frontal_only=True,
transform=transforms.Compose([ transform=transforms.Compose([
RollChannels(), # bob to skimage: RollChannels(), # bob to skimage:
transforms.ToTensor(), ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) 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) dataloader = torch.utils.data.DataLoader(face_dataset, batch_size=batch_size, shuffle=True)
...@@ -129,13 +122,13 @@ def main(user_input=None): ...@@ -129,13 +122,13 @@ def main(user_input=None):
generator.apply(weights_init) generator.apply(weights_init)
logger.info("Generator architecture: {}".format(generator)) logger.info("Generator architecture: {}".format(generator))
discrminator = DCGAN_discriminator(ngpu) discriminator = DCGAN_discriminator(ngpu)
discrminator.apply(weights_init) discriminator.apply(weights_init)
logger.info("Discriminator architecture: {}".format(discrminator)) logger.info("Discriminator architecture: {}".format(discriminator))
# =============== # ===============
# === TRAINER === # === TRAINER ===
# =============== # ===============
trainer = DCGANTrainer(generator, discrminator, batch_size=batch_size, noise_dim=noise_dim, use_gpu=use_gpu, verbosity_level=verbosity_level) trainer = DCGANTrainer(generator, discriminator, batch_size=batch_size, noise_dim=noise_dim, use_gpu=use_gpu, verbosity_level=verbosity_level)
trainer.train(dataloader, n_epochs=epochs, output_dir=output_dir) trainer.train(dataloader, n_epochs=epochs, output_dir=output_dir)
...@@ -105,7 +105,7 @@ class DCGANTrainer(object): ...@@ -105,7 +105,7 @@ class DCGANTrainer(object):
# train with real # train with real
self.netD.zero_grad() self.netD.zero_grad()
real_cpu = data real_cpu = data['image']
batch_size = real_cpu.size(0) batch_size = real_cpu.size(0)
if self.use_gpu: if self.use_gpu:
real_cpu = real_cpu.cuda() real_cpu = real_cpu.cuda()
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment