DCGAN family change

examples/GAN/BEGAN.py

@@ -131,13 +131,17 @@ class Model(GANModelDesc):
         return opt
 
 
-DCGAN.Model = Model
-
+if __name__ == '__main__':
+    args = DCGAN.get_args()
+    if args.sample:
+        DCGAN.sample(Model(), args.load, 'gen/conv4.3/output')
+    else:
+        assert args.data
+        logger.auto_set_dir()
 
-def get_config():
-    return TrainConfig(
+        config = TrainConfig(
             model=Model(),
-        dataflow=DCGAN.get_data(G.data),
+            dataflow=DCGAN.get_data(args.data),
             callbacks=[
                 ModelSaver(),
                 StatMonitorParamSetter(
@@ -145,21 +149,9 @@ def get_config():
             ],
             steps_per_epoch=500,
             max_epoch=400,
+            session_init=SaverRestore(args.load) if args.load else None,
+            nr_tower=max(get_nr_gpu(), 1)
         )
-
-
-if __name__ == '__main__':
-    args = DCGAN.get_args()
-    if args.sample:
-        DCGAN.sample(args.load, 'gen/conv4.3/output')
-    else:
-        assert args.data
-        logger.auto_set_dir()
-        config = get_config()
-        if args.load:
-            config.session_init = SaverRestore(args.load)
-        nr_gpu = get_nr_gpu()
-        config.nr_tower = max(get_nr_gpu(), 1)
         if config.nr_tower == 1:
             GANTrainer(config).train()
         else:

examples/GAN/DCGAN.py

@@ -115,7 +115,7 @@ def get_data(datadir):
     ds = ImageFromFile(imgs, channel=3, shuffle=True)
     ds = AugmentImageComponent(ds, get_augmentors())
     ds = BatchData(ds, opt.BATCH)
-    ds = PrefetchDataZMQ(ds, 1)
+    ds = PrefetchDataZMQ(ds, 5)
     return ds
 
 
@@ -129,10 +129,10 @@ def get_config():
     )
 
 
-def sample(model_path, output_name='gen/gen'):
+def sample(model, model_path, output_name='gen/gen'):
     pred = PredictConfig(
         session_init=get_model_loader(model_path),
-        model=Model(),
+        model=model,
         input_names=['z'],
         output_names=[output_name, 'z'])
     pred = SimpleDatasetPredictor(pred, RandomZData((100, opt.Z_DIM)))
@@ -162,7 +162,7 @@ def get_args():
 if __name__ == '__main__':
     args = get_args()
     if args.sample:
-        sample(args.load)
+        sample(Model(), args.load)
     else:
         assert args.data
         logger.auto_set_dir()

examples/GAN/Improved-WGAN.py

@@ -93,7 +93,7 @@ DCGAN.Model = Model
 if __name__ == '__main__':
     args = DCGAN.get_args()
     if args.sample:
-        DCGAN.sample(args.load)
+        DCGAN.sample(Model(), args.load)
     else:
         assert args.data
         logger.auto_set_dir()

examples/GAN/WGAN.py

@@ -54,7 +54,7 @@ if __name__ == '__main__':
     args = DCGAN.get_args()
 
     if args.sample:
-        DCGAN.sample(args.load)
+        DCGAN.sample(Model(), args.load)
     else:
         assert args.data
         logger.auto_set_dir()