Merge pull request #3 from SiLiKhon/refactoring

Refactoring

combine_images.py

@@ -0,0 +1,49 @@
+import argparse
+from pathlib import Path
+
+from PIL import Image
+
+
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument('path_to_images', type=str)
+  parser.add_argument('--output_name', type=str, default='plots.png')
+
+  args = parser.parse_args()
+
+  variables = [
+      'crossing_angle',
+      'dip_angle',
+      'drift_length',
+      'pad_coord_fraction',
+      'time_bin_fraction',
+  ]
+
+  stats = [
+      'Mean0',
+      'Mean1',
+      'Sigma0^2',
+      'Sigma1^2',
+      'Cov01',
+      'Sum',
+  ]
+
+  img_path = Path(args.path_to_images)
+  images = [[Image.open(img_path / f'{s} vs {v}_amp_gt_1.png') for v in variables] for s in stats]
+
+  width, height = images[0][0].size
+
+  new_image = Image.new('RGB', (width * len(stats), height * len(variables)))
+
+  x_offset = 0
+  for img_line in images:
+      y_offset = 0
+      for img in img_line:
+          new_image.paste(img, (x_offset, y_offset))
+          y_offset += img.size[1]
+      x_offset += img.size[0]
+
+  new_image.save(img_path / args.output_name)
+
+if __name__ == '__main__':
+  main()

cuda_gpu_config.py

@@ -0,0 +1,15 @@
+import os
+
+import tensorflow as tf
+
+def setup_gpu(gpu_num=None):
+    os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
+    if gpu_num is not None:
+        os.environ['CUDA_VISIBLE_DEVICES'] = 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"

dump_graph_model_v4.py

@@ -3,20 +3,26 @@
 
 import tensorflow as tf
 
+from cuda_gpu_config import setup_gpu
 from model_export import dump_graph
-from models.baseline_v4_8x16 import preprocess_features
+from models.model_v4 import preprocess_features, Model_v4
+from models.utils import load_weights
+from run_model_v4 import load_config
 
 def main():
     parser = argparse.ArgumentParser(fromfile_prefix_chars='@')
     parser.add_argument('--checkpoint_name', type=str, required=True)
     parser.add_argument('--output_path', type=str, default='model_export/model_v4/graph.pbtxt')
     parser.add_argument('--latent_dim', type=int, default=32, required=False)
-    parser.add_argument('--dont_hack_upsampling_op', default=False, action='store_true')
+    parser.add_argument('--dont_hack_upsampling_op', default=True, action='store_true')
     parser.add_argument('--test_input', type=float, nargs=4, default=None)
     parser.add_argument('--constant_seed', type=float, default=None)
+    parser.add_argument('--gpu_num', type=str, default=None)
 
     args, _ = parser.parse_known_args()
 
+    setup_gpu(args.gpu_num)
+
     print("")
     print("----" * 10)
     print("Arguments:")
@@ -30,13 +36,10 @@ def epoch_from_name(name):
         return int(epoch)
 
     model_path = Path('saved_models') / args.checkpoint_name
-    gen_checkpoints = model_path.glob("generator_*.h5")
-    latest_gen_checkpoint = max(
-        gen_checkpoints,
-        key=lambda path: epoch_from_name(path.stem)
-    )
 
-    model = tf.keras.models.load_model(str(latest_gen_checkpoint), compile=False)
+    full_model = Model_v4(load_config(model_path / 'config.yaml'))
+    load_weights(full_model, model_path)
+    model = full_model.generator
 
     if args.constant_seed is None:
         def preprocess(x):

models/baseline_fc_v4_8x16.py → legacy_code/models/baseline_fc_v4_8x16.py

@@ -1,5 +1,7 @@
 import tensorflow as tf
 
+from . import scalers
+
 @tf.function(experimental_relax_shapes=True)
 def preprocess_features(features):
     # features:
@@ -145,6 +147,10 @@ def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uni
 #                               loss='mean_squared_error')
 #        self.discriminator.compile(optimizer=self.disc_opt,
 #                               loss='mean_squared_error')
+        self.scaler = scalers.Logarithmic()
+        self.pad_range = (-3, 5)
+        self.time_range = (-7, 9)
+        self.data_version = 'data_v4'
 
 
     @tf.function