Centralize network creation for training and generating

Author: Aggrathon

generate.py

@@ -1,10 +1,7 @@
 
 import os
 
-import tensorflow as tf
-
-from network import GANetwork
-from train import CONFIG
+from train import CONFIG, get_network
 
 
 def get_config(batch):
@@ -14,20 +11,19 @@ def get_config(batch):
     config['grid_size'] = int(batch**0.5)
     return config
 
-
 def generate(name, amount=1):
-    gan = GANetwork(name, **get_config(amount))
-    session, _, iter = gan.get_session()
-    if iter == 0:
+    gan = get_network(name, **get_config(amount))
+    session, _, iteration = gan.get_session()
+    if iteration == 0:
         print("No already trained network found (%s)"%name)
         return
     print("Generating %d images using the %s network"%(amount, name))
     gan.generate(session, gan.name, amount)
 
 def generate_grid(name, size=5):
-    gan = GANetwork(name, **get_config(size*size))
-    session, _, iter = gan.get_session()
-    if iter == 0:
+    gan = get_network(name, **get_config(size*size))
+    session, _, iteration = gan.get_session()
+    if iteration == 0:
         print("No already trained network found (%s)"%name)
         return
     print("Generating a image grid using the %s network"%name)

train.py

@@ -12,12 +12,15 @@
     'generator_convolutions': 5,
 }
 
+def get_network(name, **config):
+    return GANetwork(name, **config)
+
 
 if __name__ == '__main__':
     if len(os.sys.argv) < 2:
         print('Usage:')
         print('  python %s network_name [num_iterations]\t- Trains a network on the images in the input folder'%os.sys.argv[0])
     elif len(os.sys.argv) < 3:
-        GANetwork(os.sys.argv[1], **CONFIG).train()
+        get_network(os.sys.argv[1], **CONFIG).train()
     else:
-        GANetwork(os.sys.argv[1], **CONFIG).train(int(os.sys.argv[2]))
+        get_network(os.sys.argv[1], **CONFIG).train(int(os.sys.argv[2]))