[script] generic script to train DCGAN

bob/learn/pytorch/scripts/train_dcgan_multipie.py → bob/learn/pytorch/scripts/train_dcgan.py

@@ -103,22 +103,10 @@ def main(user_input=None):
     logger.error("Please provide a dataset in your configuration file !")
     sys.exit()
   
-  # ===============
-  # === NETWORK ===
-  # ===============
-  ngpu = 1 # usually we don't have more than one GPU
-  
-  generator = DCGAN_generator(ngpu)
-  generator.apply(weights_init)
-  logger.info("Generator architecture: {}".format(generator))
-
-  discriminator = DCGAN_discriminator(ngpu)
-  discriminator.apply(weights_init)
-  logger.info("Discriminator architecture: {}".format(discriminator))
-
-
-  # ===============
-  # === TRAINER ===
-  # ===============
-  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)
+  # train the model
+  if hasattr(configuration, 'generator') and hasattr(configuration, 'discriminator'):
+    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)
+  else:
+    logger.error("Please provide both a generator and a discriminator in your configuration file !")
+    sys.exit()