[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

bob/learn/pytorch/scripts/train_dcgan_multipie.py

@@ -5,19 +5,18 @@
 """ Train a DR-GAN 
 
 Usage:
-  %(prog)s [--latent-dim=<int>] [--noise-dim=<int>] 
+  %(prog)s [--noise-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.
-  -l, --latent-dim=<int>    the dimension of the encoded ID [default: 320]
   -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: 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 
   -S, --seed=<int>          The random seed [default: 3] 
   -v, --verbose             Increase the verbosity (may appear multiple times).
@@ -26,7 +25,7 @@ Example:
 
   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.
 
@@ -53,9 +52,11 @@ import torchvision.transforms as transforms
 import torchvision.utils as vutils
 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 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_discriminator
@@ -74,7 +75,7 @@ def main(user_input=None):
 
   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,)
+  args = docopt(__doc__ % completions,argv=arguments,version='Train DCGAN (%s)' % version,)
 
   # verbosity
   verbosity_level = args['--verbose']
@@ -89,32 +90,24 @@ def main(user_input=None):
   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)
+  
+  bob.io.base.create_directories_safe(output_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=True, 
                                  transform=transforms.Compose([
                                    RollChannels(), # bob to skimage:  
-                                   transforms.ToTensor(),
-                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+                                   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)
@@ -129,13 +122,13 @@ def main(user_input=None):
   generator.apply(weights_init)
   logger.info("Generator architecture: {}".format(generator))
 
-  discrminator = DCGAN_discriminator(ngpu)
-  discrminator.apply(weights_init)
-  logger.info("Discriminator architecture: {}".format(discrminator))
+  discriminator = DCGAN_discriminator(ngpu)
+  discriminator.apply(weights_init)
+  logger.info("Discriminator architecture: {}".format(discriminator))
 
 
   # ===============
   # === 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)

bob/learn/pytorch/trainers/DCGANTrainer.py

@@ -105,7 +105,7 @@ class DCGANTrainer(object):
       
         # train with real
         self.netD.zero_grad()
-        real_cpu = data
+        real_cpu = data['image']
         batch_size = real_cpu.size(0)
         if self.use_gpu:
           real_cpu = real_cpu.cuda()