refactor for convenience

models/baseline_10x10/__init__.py

@@ -1,75 +1,86 @@
 import tensorflow as tf
-LATENT_DIM = 32
-activation = tf.keras.activations.relu
-dropout_rate = 0.2
 
+
+LATENT_DIM = 32
 NUM_DISC_UPDATES = 3
 GP_LAMBDA = 10.
 
-generator = tf.keras.Sequential([
-    tf.keras.layers.Dense(units=64, activation=activation, input_shape=(LATENT_DIM,)),
 
-    tf.keras.layers.Reshape((4, 4, 4)),
+def get_generator(activation = tf.keras.activations.relu, kernel_init='glorot_uniform'):  # he_normalZ
+    generator = tf.keras.Sequential([
+        tf.keras.layers.Dense(units=64, activation=activation, input_shape=(LATENT_DIM,)),
 
-    tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation),
-    tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation),
-    tf.keras.layers.UpSampling2D(), # 8x8
+        tf.keras.layers.Reshape((4, 4, 4)),
 
-    tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='same' , activation=activation),
-    tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='valid', activation=activation), # 6x6
-    tf.keras.layers.UpSampling2D(), # 12x12
+        tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation, kernel_initializer=kernel_init),
+        tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation, kernel_initializer=kernel_init),
+        tf.keras.layers.UpSampling2D(),  # 8x8
 
-    tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='valid', activation=activation), # 10x10
+        tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='same' , activation=activation, kernel_initializer=kernel_init),
+        tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init),  # 6x6
+        tf.keras.layers.UpSampling2D(),  # 12x12
 
-    tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu),
+        tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init),  # 10x10
 
-    tf.keras.layers.Reshape((10, 10)),
-], name='generator')
+        tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, kernel_initializer=kernel_init),
 
-discriminator = tf.keras.Sequential([
-    tf.keras.layers.Reshape((10, 10, 1), input_shape=(10, 10)),
+        tf.keras.layers.Reshape((10, 10)),
+    ], name='generator')
+    return generator
 
-    tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='same', activation=activation),
-    tf.keras.layers.Dropout(dropout_rate),
-    tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='valid', activation=activation), # 8x8
-    tf.keras.layers.Dropout(dropout_rate),
 
-    tf.keras.layers.MaxPool2D(), # 4x4
+def get_discriminator(activation=tf.keras.activations.relu, kernel_init='glorot_uniform', dropout_rate=0.2):
+    discriminator = tf.keras.Sequential([
+        tf.keras.layers.Reshape((10, 10, 1), input_shape=(10, 10)),
 
-    tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation),
-    tf.keras.layers.Dropout(dropout_rate),
-    tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation),
-    tf.keras.layers.Dropout(dropout_rate),
+        tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='same', activation=activation, kernel_initializer=kernel_init),
+        tf.keras.layers.Dropout(dropout_rate),
+        tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init),  # 8x8
+        tf.keras.layers.Dropout(dropout_rate),
 
-    tf.keras.layers.MaxPool2D(), # 2x2
+        tf.keras.layers.MaxPool2D(),  # 4x4
 
-    tf.keras.layers.Conv2D(filters=64, kernel_size=2, padding='valid', activation=activation), # 1x1
-    tf.keras.layers.Dropout(dropout_rate),
+        tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation, kernel_initializer=kernel_init),
+        tf.keras.layers.Dropout(dropout_rate),
+        tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation, kernel_initializer=kernel_init),
+        tf.keras.layers.Dropout(dropout_rate),
 
-    tf.keras.layers.Reshape((64,)),
+        tf.keras.layers.MaxPool2D(),  # 2x2
 
-    tf.keras.layers.Dense(units=128, activation=activation),
-    tf.keras.layers.Dropout(dropout_rate),
+        tf.keras.layers.Conv2D(filters=64, kernel_size=2, padding='valid', activation=activation, kernel_initializer=kernel_init),  # 1x1
+        tf.keras.layers.Dropout(dropout_rate),
 
-    tf.keras.layers.Dense(units=1, activation=None),
-], name='discriminator')
+        tf.keras.layers.Reshape((64,)),
 
+        tf.keras.layers.Dense(units=128, activation=activation),
+        tf.keras.layers.Dropout(dropout_rate),
 
+        tf.keras.layers.Dense(units=1, activation=None),
+    ], name='discriminator')
+    return discriminator
 
-disc_opt = tf.optimizers.RMSprop(0.0005)
-gen_opt  = tf.optimizers.RMSprop(0.0005)
-
-def make_fake(size):
-    return generator(
-        tf.random.normal(shape=(size, LATENT_DIM), dtype='float32')
-    )
 
 def disc_loss(d_real, d_fake):
     return tf.reduce_mean(d_fake - d_real)
 
+
 def gen_loss(d_real, d_fake):
     return tf.reduce_mean(d_real - d_fake)
 
+
+disc_opt = tf.keras.optimizers.RMSprop(1e-4)
+gen_opt  = tf.keras.optimizers.RMSprop(1e-4)
+
+generator = get_generator()
+discriminator = get_discriminator()
+
+
+def make_fake(size):
+    return generator(
+        tf.random.normal(shape=(size, LATENT_DIM), dtype='float32')
+    )
+
+
 def gradient_penalty(real, fake):
     alpha = tf.random.uniform(shape=[len(real), 1, 1])
     interpolates = alpha * real + (1 - alpha) * fake
@@ -79,6 +90,7 @@ def gradient_penalty(real, fake):
     grads = tf.reshape(t.gradient(d_int, interpolates), [len(real), -1])
     return tf.reduce_mean(tf.maximum(tf.norm(grads, axis=-1) - 1, 0)**2)
 
+
 @tf.function
 def calculate_losses(batch):
     fake = make_fake(len(batch))
@@ -89,6 +101,7 @@ def calculate_losses(batch):
     g_loss = gen_loss(d_real, d_fake)
     return {'disc_loss' : d_loss, 'gen_loss' : g_loss}
 
+
 def disc_step(batch):
     batch = tf.convert_to_tensor(batch)
 
@@ -99,6 +112,7 @@ def disc_step(batch):
     disc_opt.apply_gradients(zip(grads, discriminator.trainable_variables))
     return losses
 
+
 def gen_step(batch):
     batch = tf.convert_to_tensor(batch)
 
@@ -109,8 +123,10 @@ def gen_step(batch):
     gen_opt.apply_gradients(zip(grads, generator.trainable_variables))
     return losses
 
+
 step_counter = tf.Variable(0, dtype='int32', trainable=False)
 
+
 @tf.function
 def training_step(batch):
     if step_counter == NUM_DISC_UPDATES:

models/training.py

@@ -2,6 +2,7 @@
 import numpy as np
 from tqdm import trange
 
+
 def train(data_train, data_val, train_step_fn, loss_eval_fn, num_epochs, batch_size, train_writer=None, val_writer=None, callbacks=[]):
     for i_epoch in range(num_epochs):
         print("Working on epoch #{}".format(i_epoch), flush=True)

test_script_data_v1.py

@@ -1,67 +1,85 @@
 import os
 from pathlib import Path
-
-os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID'
-os.environ['CUDA_VISIBLE_DEVICES']='0'
-
+import argparse
 import numpy as np
 from sklearn.model_selection import train_test_split
 import tensorflow as tf
 
-gpus = tf.config.experimental.list_physical_devices('GPU')
-for gpu in gpus:
-    tf.config.experimental.set_memory_growth(gpu, True)
 
-from data import preprocessing
-from models import training, baseline_10x10
-from metrics import make_metric_plots, make_histograms
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--checkpoint_name', type=str, default='baseline_10x10', required=False)
+    parser.add_argument('--batch_size', type=int, default=32, required=False)
+    parser.add_argument('--save_every', type=int, default=50, required=False)
+    parser.add_argument('--gpu_num', type=str, default=None, required=False)
+
+    args = parser.parse_args()
+
+    os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
+    if args.gpu_num is not None:
+        os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_num
+    gpus = tf.config.experimental.list_physical_devices('GPU')
+    for gpu in gpus:
+        tf.config.experimental.set_memory_growth(gpu, True)
+
+    logical_devices = tf.config.experimental.list_logical_devices('GPU')
+    assert len(logical_devices) > 0, "Not enough GPU hardware devices available"
+
+    from data import preprocessing
+    from models import training, baseline_10x10
+    from metrics import make_metric_plots, make_histograms
+
+    model_path = Path('saved_models') / args.checkpoint_name
+    model_path.mkdir(parents=True)
 
-preprocessing._VERSION = 'data_v1'
-data = preprocessing.read_csv_2d(pad_range = (39, 49), time_range = (266, 276))
+    def save_model(step):
+        if step % args.save_every == 0:
+            print(f'Saving model on step {step} to {model_path}')
+            baseline_10x10.generator.save(str(model_path.joinpath("generator_{:05d}.h5".format(step))))
+            baseline_10x10.discriminator.save(str(model_path.joinpath("discriminator_{:05d}.h5".format(step))))
 
-data_scaled = np.log10(1 + data).astype('float32')
-X_train, X_test = train_test_split(data_scaled, test_size=0.25, random_state=42)
 
-writer_train = tf.summary.create_file_writer('logs/baseline_10x10/train')
-writer_val   = tf.summary.create_file_writer('logs/baseline_10x10/validation')
+    preprocessing._VERSION = 'data_v1'
+    data = preprocessing.read_csv_2d(pad_range=(39, 49), time_range=(266, 276))
 
-unscale = lambda x: 10**x - 1
+    data_scaled = np.log10(1 + data).astype('float32')
+    X_train, X_test = train_test_split(data_scaled, test_size=0.25, random_state=42)
 
-def write_hist_summary(step):
-    if step % 50 == 0:
-        gen_scaled = baseline_10x10.make_fake(len(X_test)).numpy()
-        real = unscale(X_test)
-        gen  = unscale(gen_scaled)
-        gen[gen < 0] = 0
-        gen1 = np.where(gen < 1., 0, gen)
-        images  = make_metric_plots(real, gen )
-        images1 = make_metric_plots(real, gen1)
+    writer_train = tf.summary.create_file_writer(f'logs/{args.checkpoint_name}/train')
+    writer_val = tf.summary.create_file_writer(f'logs/{args.checkpoint_name}/validation')
 
-        img_amplitude = make_histograms(X_test.flatten(), gen_scaled.flatten(), 'log10(amplitude + 1)', logy=True)
+    unscale = lambda x: 10 ** x - 1
 
+    def write_hist_summary(step):
+        if step % 50 == 0:
+            gen_scaled = baseline_10x10.make_fake(len(X_test)).numpy()
+            real = unscale(X_test)
+            gen = unscale(gen_scaled)
+            gen[gen < 0] = 0
+            gen1 = np.where(gen < 1., 0, gen)
+            images = make_metric_plots(real, gen)
+            images1 = make_metric_plots(real, gen1)
+
+            img_amplitude = make_histograms(X_test.flatten(), gen_scaled.flatten(), 'log10(amplitude + 1)', logy=True)
+
+            with writer_val.as_default():
+                for k, img in images.items():
+                    tf.summary.image(k, img, step)
+                for k, img in images1.items():
+                    tf.summary.image("{} (amp > 1)".format(k), img, step)
+                tf.summary.image("log10(amplitude + 1)", img_amplitude, step)
+
+    def schedule_lr(step):
+        baseline_10x10.disc_opt.lr.assign(baseline_10x10.disc_opt.lr * 0.998)
+        baseline_10x10.gen_opt.lr.assign(baseline_10x10.gen_opt.lr * 0.998)
         with writer_val.as_default():
-            for k, img in images.items():
-                tf.summary.image(k, img, step)
-            for k, img in images1.items():
-                tf.summary.image("{} (amp > 1)".format(k), img, step)
-            tf.summary.image("log10(amplitude + 1)", img_amplitude, step)
-
-model_path = Path("saved_models/baseline_10x10/")
-model_path.mkdir(parents=True)
-
-def save_model(step):
-    if step % 50 == 0:
-        baseline_10x10.generator    .save(str(model_path.joinpath("generator_{:05d}.h5"    .format(step))))
-        baseline_10x10.discriminator.save(str(model_path.joinpath("discriminator_{:05d}.h5".format(step))))
-
-def schedule_lr(step):
-    baseline_10x10.disc_opt.lr.assign(baseline_10x10.disc_opt.lr * 0.998)
-    baseline_10x10.gen_opt .lr.assign(baseline_10x10.gen_opt .lr * 0.998)
-    with writer_val.as_default():
-        tf.summary.scalar("discriminator learning rate", baseline_10x10.disc_opt.lr, step)
-        tf.summary.scalar("generator learning rate"    , baseline_10x10.gen_opt .lr, step)
-
-
-training.train(X_train, X_test, baseline_10x10.training_step, baseline_10x10.calculate_losses, 10000, 32,
-               train_writer=writer_train, val_writer=writer_val,
-               callbacks=[write_hist_summary, save_model, schedule_lr])
+            tf.summary.scalar("discriminator learning rate", baseline_10x10.disc_opt.lr, step)
+            tf.summary.scalar("generator learning rate", baseline_10x10.gen_opt.lr, step)
+
+    training.train(X_train, X_test, baseline_10x10.training_step, baseline_10x10.calculate_losses, 10000, args.batch_size,
+                   train_writer=writer_train, val_writer=writer_val,
+                   callbacks=[write_hist_summary, save_model, schedule_lr])
+
+
+if __name__ == '__main__':
+    main()