From 75da9ca4a5619240203055376fe3e663c81363b6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A1=D1=83=D1=85=D0=BE=D1=80=D0=BE=D1=81=D0=BE=D0=B2=20?= =?UTF-8?q?=D0=90=D0=BB=D0=B5=D0=BA=D1=81=D0=B5=D0=B9?= Date: Wed, 6 Oct 2021 21:52:37 +0300 Subject: [PATCH 1/6] flake8 action --- .github/workflows/push.yml | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 .github/workflows/push.yml diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml new file mode 100644 index 0000000..8ae7115 --- /dev/null +++ b/.github/workflows/push.yml @@ -0,0 +1,9 @@ +name: Flake8 + pytest stylecheck +on: [push] +jobs: + flake8: + runs-on: ubuntu-latest + steps: + - name: Run flake8 + image: docker/Dockerfile + From 74153448bb6bed6a3cfb46c21ee663fa4cebd046 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A1=D1=83=D1=85=D0=BE=D1=80=D0=BE=D1=81=D0=BE=D0=B2=20?= =?UTF-8?q?=D0=90=D0=BB=D0=B5=D0=BA=D1=81=D0=B5=D0=B9?= Date: Wed, 6 Oct 2021 22:45:58 +0300 Subject: [PATCH 2/6] style fixes --- .flake8 | 7 ++ .github/workflows/push.yml | 12 +++- combine_images.py | 81 ++++++++++++----------- cuda_gpu_config.py | 1 + data/preprocessing.py | 32 ++++----- dump_graph_model_v4.py | 6 +- legacy_code/models/baseline_10x10.py | 39 +++++++---- legacy_code/models/baseline_10x15.py | 22 ++++-- legacy_code/models/baseline_fc_v4_8x16.py | 30 ++++++--- legacy_code/models/baseline_v2_10x10.py | 37 +++++++---- legacy_code/models/baseline_v3_6x15.py | 46 ++++++++----- legacy_code/models/baseline_v4_8x16.py | 54 ++++++++++----- legacy_code/test_script_data_v0.py | 20 ++++-- legacy_code/test_script_data_v1.py | 3 +- legacy_code/test_script_data_v1_normed.py | 28 +++++--- legacy_code/test_script_data_v2.py | 32 +++++---- legacy_code/test_script_data_v3.py | 37 +++++------ legacy_code/test_script_data_v4.py | 28 ++++---- metrics/__init__.py | 55 +++++++-------- metrics/gaussian_metrics.py | 12 ++-- metrics/plotting.py | 22 +++--- metrics/trends.py | 18 ++--- model_export/dump_graph.py | 12 ++-- models/callbacks.py | 3 +- models/model_v4.py | 14 ++-- models/nn.py | 5 +- models/scalers.py | 16 ++--- models/training.py | 11 ++- models/utils.py | 4 +- run_model_v4.py | 22 +++--- 30 files changed, 408 insertions(+), 301 deletions(-) create mode 100644 .flake8 diff --git a/.flake8 b/.flake8 new file mode 100644 index 0000000..9fd30a6 --- /dev/null +++ b/.flake8 @@ -0,0 +1,7 @@ +[flake8] +max-line-length = 120 +exclude = + model_export/tf2xla_pb2.py +per-file-ignores = + legacy_code/test_script_data_v0.py:E402 + legacy_code/test_script_data_v1_normed.py:E402 \ No newline at end of file diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 8ae7115..62db9e7 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -1,9 +1,15 @@ -name: Flake8 + pytest stylecheck +name: Flake8 on: [push] jobs: flake8: runs-on: ubuntu-latest steps: + - uses: actions/checkout@v2 + - name: Install Python 3 + uses: actions/setup-python@v1 + with: + python-version: 3.9 + - name: Install flake8 + run: python -m pip install flake8 - name: Run flake8 - image: docker/Dockerfile - + run: flake8 diff --git a/combine_images.py b/combine_images.py index 622a207..41d0dc1 100644 --- a/combine_images.py +++ b/combine_images.py @@ -5,45 +5,46 @@ 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) + 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() + main() diff --git a/cuda_gpu_config.py b/cuda_gpu_config.py index 73e860d..deeaefb 100644 --- a/cuda_gpu_config.py +++ b/cuda_gpu_config.py @@ -2,6 +2,7 @@ import tensorflow as tf + def setup_gpu(gpu_num=None): os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID' if gpu_num is not None: diff --git a/data/preprocessing.py b/data/preprocessing.py index ab64c7e..a29c2c3 100644 --- a/data/preprocessing.py +++ b/data/preprocessing.py @@ -7,6 +7,7 @@ _THIS_PATH = Path(os.path.realpath(__file__)).parent _VERSION = 'data_v1' + class Reader: def __init__(self, variables, types): assert len(variables) == len(types), 'Reader.__init__: variables and types have different length' @@ -25,7 +26,6 @@ def build(self): return pd.DataFrame(self.data, columns=['evtId'] + self.vars).set_index('evtId') - def raw_to_csv(fname_in=None, fname_out=None): if fname_in is None: fname_in = str(_THIS_PATH.joinpath(_VERSION, 'raw', 'digits.dat')) @@ -39,8 +39,8 @@ def raw_to_csv(fname_in=None, fname_out=None): lines = f.readlines() reader_main = Reader( - variables = ['ipad', 'itime', 'amp'], - types = [int , int , float] + variables=['ipad', 'itime', 'amp'], + types=[int, int, float] ) data_sources = [lines] @@ -51,24 +51,24 @@ def raw_to_csv(fname_in=None, fname_out=None): if _VERSION == 'data_v2': reader_features = Reader( - variables = ["crossing_angle", "dip_angle"], - types = [float , float ] + variables=["crossing_angle", "dip_angle"], + types=[float, float] ) elif _VERSION == 'data_v3': reader_features = Reader( - variables = ["crossing_angle", "dip_angle", "drift_length"], - types = [float , float , float ] + variables=["crossing_angle", "dip_angle", "drift_length"], + types=[float, float, float] ) elif _VERSION == 'data_v4': reader_features = Reader( - variables = ["crossing_angle", "dip_angle", "drift_length", "pad_coordinate"], - types = [float , float , float , float ] + variables=["crossing_angle", "dip_angle", "drift_length", "pad_coordinate"], + types=[float, float, float, float] ) else: raise NotImplementedError lines, lines_angles = lines[1::2], lines[::2] - lines_angles = [' '.join(l.split()[1:]) for l in lines_angles] + lines_angles = [' '.join(line.split()[1:]) for line in lines_angles] data_sources = [lines, lines_angles] readers = [reader_main, reader_features] @@ -76,17 +76,19 @@ def raw_to_csv(fname_in=None, fname_out=None): for evt_id, lines_tuple in enumerate(zip(*data_sources)): for r, l in zip(readers, lines_tuple): r.read_line(l, evt_id) - + result = pd.concat([r.build() for r in readers], axis=1).reset_index() result.to_csv(fname_out, index=False) + def read_csv_2d(filename=None, pad_range=(40, 50), time_range=(265, 280), strict=True): if filename is None: filename = str(_THIS_PATH.joinpath(_VERSION, 'csv', 'digits.csv')) df = pd.read_csv(filename) - sel = lambda df, col, limits: (df[col] >= limits[0]) & (df[col] < limits[1]) + def sel(df, col, limits): + return (df[col] >= limits[0]) & (df[col] < limits[1]) if 'drift_length' in df.columns: df['itime'] -= df['drift_length'].astype(int) @@ -96,7 +98,7 @@ def read_csv_2d(filename=None, pad_range=(40, 50), time_range=(265, 280), strict selection = ( sel(df, 'itime', time_range) & - sel(df, 'ipad' , pad_range ) + sel(df, 'ipad', pad_range) ) if not selection.all(): @@ -107,9 +109,9 @@ def read_csv_2d(filename=None, pad_range=(40, 50), time_range=(265, 280), strict g = df[selection].groupby('evtId') def convert_event(event): - result = np.zeros(dtype=float, shape=(pad_range [1] - pad_range [0], + result = np.zeros(dtype=float, shape=(pad_range[1] - pad_range[0], time_range[1] - time_range[0])) - + indices = tuple(event[['ipad', 'itime']].values.T - np.array([[pad_range[0]], [time_range[0]]])) result[indices] = event.amp.values diff --git a/dump_graph_model_v4.py b/dump_graph_model_v4.py index 0889bf0..2f7483c 100644 --- a/dump_graph_model_v4.py +++ b/dump_graph_model_v4.py @@ -1,4 +1,5 @@ -import argparse, re +import argparse +import re from pathlib import Path import tensorflow as tf @@ -9,6 +10,7 @@ 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) @@ -32,7 +34,7 @@ def main(): print("") def epoch_from_name(name): - epoch, = re.findall('\d+', name) + epoch, = re.findall(r'\d+', name) return int(epoch) model_path = Path('saved_models') / args.checkpoint_name diff --git a/legacy_code/models/baseline_10x10.py b/legacy_code/models/baseline_10x10.py index 402a134..0c64cf5 100644 --- a/legacy_code/models/baseline_10x10.py +++ b/legacy_code/models/baseline_10x10.py @@ -7,16 +7,22 @@ def get_generator(activation, kernel_init, latent_dim): tf.keras.layers.Reshape((4, 4, 4)), - 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.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=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.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=8, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init), # 10x10 - tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, kernel_initializer=kernel_init), + tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='valid', activation=activation, + kernel_initializer=kernel_init), # 10x10 + tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, + kernel_initializer=kernel_init), tf.keras.layers.Reshape((10, 10)), ], name='generator') @@ -27,21 +33,26 @@ def get_discriminator(activation, kernel_init, dropout_rate): discriminator = tf.keras.Sequential([ tf.keras.layers.Reshape((10, 10, 1), input_shape=(10, 10)), - 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='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.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(), # 4x4 - 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.Dropout(dropout_rate), - 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.Dropout(dropout_rate), tf.keras.layers.MaxPool2D(), # 2x2 - tf.keras.layers.Conv2D(filters=64, kernel_size=2, padding='valid', activation=activation, kernel_initializer=kernel_init), # 1x1 + 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.Reshape((64,)), @@ -72,7 +83,11 @@ def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uni self.num_disc_updates = num_disc_updates self.generator = get_generator(activation=activation, kernel_init=kernel_init, latent_dim=latent_dim) - self.discriminator = get_discriminator(activation=activation, kernel_init=kernel_init, dropout_rate=dropout_rate) + self.discriminator = get_discriminator( + activation=activation, + kernel_init=kernel_init, + dropout_rate=dropout_rate + ) self.step_counter = tf.Variable(0, dtype='int32', trainable=False) diff --git a/legacy_code/models/baseline_10x15.py b/legacy_code/models/baseline_10x15.py index a90961d..b2099af 100644 --- a/legacy_code/models/baseline_10x15.py +++ b/legacy_code/models/baseline_10x15.py @@ -13,10 +13,10 @@ tf.keras.layers.Reshape((3, 4, 40)), tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation), - tf.keras.layers.UpSampling2D(), # 6x8 + tf.keras.layers.UpSampling2D(), # 6x8 tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='same', activation=activation), - tf.keras.layers.UpSampling2D(), # 12x16 + tf.keras.layers.UpSampling2D(), # 12x16 tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='same', activation=activation), tf.keras.layers.Conv2D(filters=4, kernel_size=(3, 2), padding='valid', activation=activation), @@ -32,14 +32,14 @@ tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='same', activation=activation), tf.keras.layers.Dropout(dropout_rate), - tf.keras.layers.MaxPool2D(padding='same'), # 5x8 + tf.keras.layers.MaxPool2D(padding='same'), # 5x8 tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation), tf.keras.layers.Dropout(dropout_rate), - tf.keras.layers.MaxPool2D(padding='same'), # 3x4 + tf.keras.layers.MaxPool2D(padding='same'), # 3x4 - tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation), # 1x2 + tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation), # 1x2 tf.keras.layers.Dropout(dropout_rate), tf.keras.layers.Reshape((128,)), @@ -51,21 +51,24 @@ ], name='discriminator') - disc_opt = tf.optimizers.RMSprop() gen_opt = tf.optimizers.RMSprop() + 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) + def gradient_penalty(real, fake): alpha = tf.random.uniform(shape=[len(real), 1, 1]) interpolates = alpha * real + (1 - alpha) * fake @@ -75,6 +78,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)) @@ -83,7 +87,8 @@ def calculate_losses(batch): d_loss = disc_loss(d_real, d_fake) + GP_LAMBDA * gradient_penalty(batch, fake) g_loss = gen_loss(d_real, d_fake) - return {'disc_loss' : d_loss, 'gen_loss' : g_loss} + return {'disc_loss': d_loss, 'gen_loss': g_loss} + def disc_step(batch): batch = tf.convert_to_tensor(batch) @@ -95,6 +100,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) @@ -105,8 +111,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: diff --git a/legacy_code/models/baseline_fc_v4_8x16.py b/legacy_code/models/baseline_fc_v4_8x16.py index c0339d5..d466e8e 100644 --- a/legacy_code/models/baseline_fc_v4_8x16.py +++ b/legacy_code/models/baseline_fc_v4_8x16.py @@ -2,6 +2,7 @@ from . import scalers + @tf.function(experimental_relax_shapes=True) def preprocess_features(features): # features: @@ -9,14 +10,16 @@ def preprocess_features(features): # dip_angle [-60, 60] # drift_length [35, 290] # pad_coordinate [40-something, 40-something] - bin_fractions = features[:,-2:] % 1 + bin_fractions = features[:, -2:] % 1 features = ( - features[:,:3] - tf.constant([[0., 0., 162.5]]) + features[:, :3] - tf.constant([[0., 0., 162.5]]) ) / tf.constant([[20., 60., 127.5]]) return tf.concat([features, bin_fractions], axis=-1) + _f = preprocess_features + def get_generator(activation, kernel_init, num_features, latent_dim): generator = tf.keras.Sequential([ tf.keras.layers.Dense(units=32, activation=activation, input_shape=(num_features + latent_dim,)), @@ -46,23 +49,29 @@ def get_discriminator(activation, kernel_init, dropout_rate, num_features, num_a ) discriminator_tail = tf.keras.Sequential([ - 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='same', activation=activation, + kernel_initializer=kernel_init), 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.Conv2D(filters=16, kernel_size=3, padding='same', activation=activation, + kernel_initializer=kernel_init), tf.keras.layers.Dropout(dropout_rate), tf.keras.layers.MaxPool2D(pool_size=(1, 2)), # 8x8 - 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.Dropout(dropout_rate), - 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.Dropout(dropout_rate), tf.keras.layers.MaxPool2D(), # 4x4 - tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init), # 2x2 + tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation, + kernel_initializer=kernel_init), # 2x2 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.Conv2D(filters=64, kernel_size=2, padding='valid', activation=activation, + kernel_initializer=kernel_init), # 1x1 tf.keras.layers.Dropout(dropout_rate), tf.keras.layers.Reshape((64,)) @@ -109,14 +118,16 @@ def gen_loss(d_real, d_fake): def disc_loss_cramer(d_real, d_fake, d_fake_2): return -tf.reduce_mean( tf.norm(d_real - d_fake, axis=-1) + - tf.norm(d_fake_2, axis=-1) - + tf.norm(d_fake_2, axis=-1) - tf.norm(d_fake - d_fake_2, axis=-1) - tf.norm(d_real, axis=-1) ) + def gen_loss_cramer(d_real, d_fake, d_fake_2): return -disc_loss_cramer(d_real, d_fake, d_fake_2) + class BaselineModel_8x16: def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uniform', dropout_rate=0.02, lr=1e-4, latent_dim=32, gp_lambda=10., num_disc_updates=8, @@ -152,7 +163,6 @@ def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uni self.time_range = (-7, 9) self.data_version = 'data_v4' - @tf.function def make_fake(self, features): size = tf.shape(features)[0] diff --git a/legacy_code/models/baseline_v2_10x10.py b/legacy_code/models/baseline_v2_10x10.py index 74d4d1c..3f95e0d 100644 --- a/legacy_code/models/baseline_v2_10x10.py +++ b/legacy_code/models/baseline_v2_10x10.py @@ -7,16 +7,22 @@ def get_generator(activation, kernel_init, num_features, latent_dim): tf.keras.layers.Reshape((4, 4, 4)), - 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.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=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.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=8, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init), # 10x10 - tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, kernel_initializer=kernel_init), + tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='valid', activation=activation, + kernel_initializer=kernel_init), # 10x10 + tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, + kernel_initializer=kernel_init), tf.keras.layers.Reshape((10, 10)), ], name='generator') @@ -40,21 +46,26 @@ def get_discriminator(activation, kernel_init, dropout_rate, num_features, num_a ) discriminator_tail = tf.keras.Sequential([ - 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='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.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(), # 4x4 - 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.Dropout(dropout_rate), - 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.Dropout(dropout_rate), tf.keras.layers.MaxPool2D(), # 2x2 - tf.keras.layers.Conv2D(filters=64, kernel_size=2, padding='valid', activation=activation, kernel_initializer=kernel_init), # 1x1 + 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.Reshape((64,)) @@ -101,14 +112,16 @@ def gen_loss(d_real, d_fake): def disc_loss_cramer(d_real, d_fake, d_fake_2): return -tf.reduce_mean( tf.norm(d_real - d_fake, axis=-1) + - tf.norm(d_fake_2, axis=-1) - + tf.norm(d_fake_2, axis=-1) - tf.norm(d_fake - d_fake_2, axis=-1) - tf.norm(d_real, axis=-1) ) + def gen_loss_cramer(d_real, d_fake, d_fake_2): return -disc_loss_cramer(d_real, d_fake, d_fake_2) + class BaselineModel10x10: def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uniform', dropout_rate=0.2, lr=1e-4, latent_dim=32, gp_lambda=10., num_disc_updates=3, diff --git a/legacy_code/models/baseline_v3_6x15.py b/legacy_code/models/baseline_v3_6x15.py index 356d6f5..6548dfc 100644 --- a/legacy_code/models/baseline_v3_6x15.py +++ b/legacy_code/models/baseline_v3_6x15.py @@ -1,36 +1,46 @@ import tensorflow as tf + @tf.function(experimental_relax_shapes=True) def preprocess_features(features): # features: # crossing_angle [-20, 20] # dip_angle [-60, 60] # drift_length [35, 290] - bin_fraction = features[:,-1:] % 1 + bin_fraction = features[:, -1:] % 1 features = ( features - tf.constant([[0., 0., 162.5]]) ) / tf.constant([[20., 60., 127.5]]) return tf.concat([features, bin_fraction], axis=-1) + _f = preprocess_features + def get_generator(activation, kernel_init, num_features, latent_dim): generator = tf.keras.Sequential([ tf.keras.layers.Dense(units=64, activation=activation, input_shape=(num_features + latent_dim,)), tf.keras.layers.Reshape((4, 4, 4)), - 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.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(size=(1, 2)), # 4x8 - 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='same', activation=activation, kernel_initializer=kernel_init), + 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='same', activation=activation, + kernel_initializer=kernel_init), tf.keras.layers.UpSampling2D(), # 8x16 - tf.keras.layers.Conv2D(filters=8, kernel_size=(1, 3), padding='same' , activation=activation, kernel_initializer=kernel_init), - tf.keras.layers.Conv2D(filters=8, kernel_size=(3, 2), padding='valid', activation=activation, kernel_initializer=kernel_init), # 6x15 - tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, kernel_initializer=kernel_init), + tf.keras.layers.Conv2D(filters=8, kernel_size=(1, 3), padding='same', activation=activation, + kernel_initializer=kernel_init), + tf.keras.layers.Conv2D(filters=8, kernel_size=(3, 2), padding='valid', activation=activation, + kernel_initializer=kernel_init), # 6x15 + tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, + kernel_initializer=kernel_init), tf.keras.layers.Reshape((6, 15)), ], name='generator') @@ -54,21 +64,26 @@ def get_discriminator(activation, kernel_init, dropout_rate, num_features, num_a ) discriminator_tail = tf.keras.Sequential([ - 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='same', activation=activation, + kernel_initializer=kernel_init), tf.keras.layers.Dropout(dropout_rate), - tf.keras.layers.Conv2D(filters=16, kernel_size=(1, 4), padding='valid', activation=activation, kernel_initializer=kernel_init), # 6x12 + tf.keras.layers.Conv2D(filters=16, kernel_size=(1, 4), padding='valid', activation=activation, + kernel_initializer=kernel_init), # 6x12 tf.keras.layers.Dropout(dropout_rate), tf.keras.layers.MaxPool2D(pool_size=(1, 2)), # 6x6 - 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.Dropout(dropout_rate), - 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.Dropout(dropout_rate), tf.keras.layers.MaxPool2D(), # 3x3 - tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init), # 1x1 + tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation, + kernel_initializer=kernel_init), # 1x1 tf.keras.layers.Dropout(dropout_rate), tf.keras.layers.Reshape((64,)) @@ -115,14 +130,16 @@ def gen_loss(d_real, d_fake): def disc_loss_cramer(d_real, d_fake, d_fake_2): return -tf.reduce_mean( tf.norm(d_real - d_fake, axis=-1) + - tf.norm(d_fake_2, axis=-1) - + tf.norm(d_fake_2, axis=-1) - tf.norm(d_fake - d_fake_2, axis=-1) - tf.norm(d_real, axis=-1) ) + def gen_loss_cramer(d_real, d_fake, d_fake_2): return -disc_loss_cramer(d_real, d_fake, d_fake_2) + class BaselineModel_6x15: def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uniform', dropout_rate=0.02, lr=1e-4, latent_dim=32, gp_lambda=10., num_disc_updates=8, @@ -154,7 +171,6 @@ def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uni # self.discriminator.compile(optimizer=self.disc_opt, # loss='mean_squared_error') - @tf.function def make_fake(self, features): size = tf.shape(features)[0] diff --git a/legacy_code/models/baseline_v4_8x16.py b/legacy_code/models/baseline_v4_8x16.py index 0511c30..c992ea2 100644 --- a/legacy_code/models/baseline_v4_8x16.py +++ b/legacy_code/models/baseline_v4_8x16.py @@ -1,5 +1,6 @@ import tensorflow as tf + @tf.function(experimental_relax_shapes=True) def preprocess_features(features): # features: @@ -7,34 +8,44 @@ def preprocess_features(features): # dip_angle [-60, 60] # drift_length [35, 290] # pad_coordinate [40-something, 40-something] - bin_fractions = features[:,-2:] % 1 + bin_fractions = features[:, -2:] % 1 features = ( - features[:,:3] - tf.constant([[0., 0., 162.5]]) + features[:, :3] - tf.constant([[0., 0., 162.5]]) ) / tf.constant([[20., 60., 127.5]]) return tf.concat([features, bin_fractions], axis=-1) + _f = preprocess_features + def get_generator(activation, kernel_init, num_features, latent_dim): generator = tf.keras.Sequential([ tf.keras.layers.Dense(units=64, activation=activation, input_shape=(num_features + latent_dim,)), tf.keras.layers.Reshape((2, 4, 8)), - tf.keras.layers.Conv2D(filters=32, kernel_size=2, padding='same', activation=activation, kernel_initializer=kernel_init), - tf.keras.layers.Conv2D(filters=32, kernel_size=2, padding='same', activation=activation, kernel_initializer=kernel_init), # 2x4 + tf.keras.layers.Conv2D(filters=32, kernel_size=2, padding='same', activation=activation, + kernel_initializer=kernel_init), + tf.keras.layers.Conv2D(filters=32, kernel_size=2, padding='same', activation=activation, + kernel_initializer=kernel_init), # 2x4 tf.keras.layers.UpSampling2D(), # 4x8 - tf.keras.layers.Conv2D(filters=16, kernel_size=(2, 3), padding='valid', activation=activation, kernel_initializer=kernel_init), # 3x6 - tf.keras.layers.Conv2D(filters=16, kernel_size=(1, 2), padding='valid', activation=activation, kernel_initializer=kernel_init), # 3x5 + tf.keras.layers.Conv2D(filters=16, kernel_size=(2, 3), padding='valid', activation=activation, + kernel_initializer=kernel_init), # 3x6 + tf.keras.layers.Conv2D(filters=16, kernel_size=(1, 2), padding='valid', activation=activation, + kernel_initializer=kernel_init), # 3x5 tf.keras.layers.UpSampling2D(), # 6x10 - tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='same' , activation=activation, kernel_initializer=kernel_init), - tf.keras.layers.Conv2D(filters=8, kernel_size=2, padding='valid', activation=activation, kernel_initializer=kernel_init), # 5x9 + tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='same', activation=activation, + kernel_initializer=kernel_init), + tf.keras.layers.Conv2D(filters=8, kernel_size=2, padding='valid', activation=activation, + kernel_initializer=kernel_init), # 5x9 tf.keras.layers.UpSampling2D(), # 10x18 - tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init), # 8x16 - tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, kernel_initializer=kernel_init), + tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='valid', activation=activation, + kernel_initializer=kernel_init), # 8x16 + tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, + kernel_initializer=kernel_init), tf.keras.layers.Reshape((8, 16)), ], name='generator') @@ -58,23 +69,29 @@ def get_discriminator(activation, kernel_init, dropout_rate, num_features, num_a ) discriminator_tail = tf.keras.Sequential([ - 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='same', activation=activation, + kernel_initializer=kernel_init), 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.Conv2D(filters=16, kernel_size=3, padding='same', activation=activation, + kernel_initializer=kernel_init), tf.keras.layers.Dropout(dropout_rate), tf.keras.layers.MaxPool2D(pool_size=(1, 2)), # 8x8 - 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.Dropout(dropout_rate), - 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.Dropout(dropout_rate), tf.keras.layers.MaxPool2D(), # 4x4 - tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init), # 2x2 + tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation, + kernel_initializer=kernel_init), # 2x2 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.Conv2D(filters=64, kernel_size=2, padding='valid', activation=activation, + kernel_initializer=kernel_init), # 1x1 tf.keras.layers.Dropout(dropout_rate), tf.keras.layers.Reshape((64,)) @@ -121,14 +138,16 @@ def gen_loss(d_real, d_fake): def disc_loss_cramer(d_real, d_fake, d_fake_2): return -tf.reduce_mean( tf.norm(d_real - d_fake, axis=-1) + - tf.norm(d_fake_2, axis=-1) - + tf.norm(d_fake_2, axis=-1) - tf.norm(d_fake - d_fake_2, axis=-1) - tf.norm(d_real, axis=-1) ) + def gen_loss_cramer(d_real, d_fake, d_fake_2): return -disc_loss_cramer(d_real, d_fake, d_fake_2) + class BaselineModel_8x16: def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uniform', dropout_rate=0.02, lr=1e-4, latent_dim=32, gp_lambda=10., num_disc_updates=8, @@ -160,7 +179,6 @@ def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uni # self.discriminator.compile(optimizer=self.disc_opt, # loss='mean_squared_error') - @tf.function def make_fake(self, features): size = tf.shape(features)[0] diff --git a/legacy_code/test_script_data_v0.py b/legacy_code/test_script_data_v0.py index f3a5a63..c0cbf0c 100644 --- a/legacy_code/test_script_data_v0.py +++ b/legacy_code/test_script_data_v0.py @@ -1,8 +1,8 @@ import os from pathlib import Path -os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID' -os.environ['CUDA_VISIBLE_DEVICES']='1' +os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID' +os.environ['CUDA_VISIBLE_DEVICES'] = '1' import numpy as np from sklearn.model_selection import train_test_split @@ -23,18 +23,21 @@ 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_10x15/train') -writer_val = tf.summary.create_file_writer('logs/baseline_10x15/validation') +writer_val = tf.summary.create_file_writer('logs/baseline_10x15/validation') + + +def unscale(x): + return 10 ** x - 1 -unscale = lambda x: 10**x - 1 def write_hist_summary(step): if step % 50 == 0: gen_scaled = baseline_10x15.make_fake(len(X_test)).numpy() real = unscale(X_test) - gen = unscale(gen_scaled) + gen = unscale(gen_scaled) gen[gen < 0] = 0 gen30 = np.where(gen < 30., 0, gen) - images = make_metric_plots(real, gen) + images = make_metric_plots(real, gen) images30 = make_metric_plots(real, gen30) img_amplitude = make_histograms(X_test.flatten(), gen_scaled.flatten(), 'log10(amplitude + 1)', logy=True) @@ -46,20 +49,23 @@ def write_hist_summary(step): tf.summary.image("{} (amp > 30)".format(k), img, step) tf.summary.image("log10(amplitude + 1)", img_amplitude, step) + model_path = Path("saved_models/baseline_10x15/") model_path.mkdir(parents=True) + def save_model(step): if step % 50 == 0: baseline_10x15.generator .save(str(model_path.joinpath("generator_{:05d}.h5" .format(step)))) baseline_10x15.discriminator.save(str(model_path.joinpath("discriminator_{:05d}.h5".format(step)))) + def schedule_lr(step): baseline_10x15.disc_opt.lr.assign(baseline_10x15.disc_opt.lr * 0.998) baseline_10x15.gen_opt .lr.assign(baseline_10x15.gen_opt .lr * 0.998) with writer_val.as_default(): tf.summary.scalar("discriminator learning rate", baseline_10x15.disc_opt.lr, step) - tf.summary.scalar("generator learning rate" , baseline_10x15.gen_opt .lr, step) + tf.summary.scalar("generator learning rate", baseline_10x15.gen_opt.lr, step) training.train(X_train, X_test, baseline_10x15.training_step, baseline_10x15.calculate_losses, 10000, 32, diff --git a/legacy_code/test_script_data_v1.py b/legacy_code/test_script_data_v1.py index f75f94d..4f9e9e7 100644 --- a/legacy_code/test_script_data_v1.py +++ b/legacy_code/test_script_data_v1.py @@ -56,7 +56,8 @@ def save_model(step): 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') - unscale = lambda x: 10 ** x - 1 + def unscale(x): + return 10 ** x - 1 def write_hist_summary(step): if step % args.save_every == 0: diff --git a/legacy_code/test_script_data_v1_normed.py b/legacy_code/test_script_data_v1_normed.py index 21f6366..ca14403 100644 --- a/legacy_code/test_script_data_v1_normed.py +++ b/legacy_code/test_script_data_v1_normed.py @@ -1,10 +1,10 @@ import os from pathlib import Path -os.environ['CUDA_DEVICE_ORDER']='PCI_BUS_ID' -os.environ['CUDA_VISIBLE_DEVICES']='0' +os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID' +os.environ['CUDA_VISIBLE_DEVICES'] = '0' -import numpy as np +import numpy as np # from sklearn.model_selection import train_test_split import tensorflow as tf @@ -14,10 +14,10 @@ from data import preprocessing from models import training, baseline_10x10 as baseline_10x10_normed -from metrics import make_metric_plots, make_histograms +from metrics import make_metric_plots preprocessing._VERSION = 'data_v1' -data = preprocessing.read_csv_2d(pad_range = (39, 49), time_range = (266, 276)) +data = preprocessing.read_csv_2d(pad_range=(39, 49), time_range=(266, 276)) data /= data.sum(axis=(1, 2), keepdims=True) / 100 data_scaled = np.log10(1 + data).astype('float32') @@ -27,36 +27,42 @@ 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_normed/train') -writer_val = tf.summary.create_file_writer('logs/baseline_10x10_normed/validation') +writer_val = tf.summary.create_file_writer('logs/baseline_10x10_normed/validation') + + +def unscale(x): + return 10 ** x - 1 -unscale = lambda x: 10**x - 1 def write_hist_summary(step): if step % 50 == 0: gen_scaled = baseline_10x10_normed.make_fake(len(X_test)).numpy() real = unscale(X_test) - gen = unscale(gen_scaled) + gen = unscale(gen_scaled) gen[gen < 0] = 0 - images = make_metric_plots(real, gen ) + images = make_metric_plots(real, gen) with writer_val.as_default(): for k, img in images.items(): tf.summary.image(k, img, step) + model_path = Path("saved_models/baseline_10x10_normed/") model_path.mkdir(parents=True) + def save_model(step): if step % 50 == 0: - baseline_10x10_normed.generator .save(str(model_path.joinpath("generator_{:05d}.h5" .format(step)))) + baseline_10x10_normed.generator.save(str(model_path.joinpath("generator_{:05d}.h5".format(step)))) baseline_10x10_normed.discriminator.save(str(model_path.joinpath("discriminator_{:05d}.h5".format(step)))) + def schedule_lr(step): baseline_10x10_normed.disc_opt.lr.assign(baseline_10x10_normed.disc_opt.lr * 0.998) baseline_10x10_normed.gen_opt .lr.assign(baseline_10x10_normed.gen_opt .lr * 0.998) with writer_val.as_default(): tf.summary.scalar("discriminator learning rate", baseline_10x10_normed.disc_opt.lr, step) - tf.summary.scalar("generator learning rate" , baseline_10x10_normed.gen_opt .lr, step) + tf.summary.scalar("generator learning rate", baseline_10x10_normed.gen_opt .lr, step) baseline_10x10_normed.disc_opt.lr.assign(0.0001) diff --git a/legacy_code/test_script_data_v2.py b/legacy_code/test_script_data_v2.py index 402fca5..836d184 100644 --- a/legacy_code/test_script_data_v2.py +++ b/legacy_code/test_script_data_v2.py @@ -1,4 +1,5 @@ -import os, sys +import os +import sys import re from pathlib import Path import argparse @@ -38,10 +39,8 @@ def main(): parser.add_argument('--feature_noise_power', type=float, default=None) parser.add_argument('--feature_noise_decay', type=float, default=None) - args = parser.parse_args() - assert ( (args.feature_noise_power is None) == (args.feature_noise_decay is None) @@ -51,7 +50,7 @@ def main(): print("----" * 10) print("Arguments:") for k, v in vars(args).items(): - print(f" {k} : {v}") + print(f" {k}: {v}") print("----" * 10) print("") @@ -79,7 +78,8 @@ def main(): f.write('\n'.join(raw_args)) for fname in fnames: with open(fname, 'r') as f_in: - if len(raw_args) > 0: f.write('\n') + if len(raw_args) > 0: + f.write('\n') f.write(f_in.read()) model = BaselineModel10x10(kernel_init=args.kernel_init, lr=args.lr, @@ -92,7 +92,7 @@ def main(): if args.prediction_only: def epoch_from_name(name): - epoch, = re.findall('\d+', name) + epoch, = re.findall(r'\d+', name) return int(epoch) gen_checkpoints = model_path.glob("generator_*.h5") @@ -115,8 +115,6 @@ def epoch_from_name(name): print(f'Loading discriminator weights from {str(latest_disc_checkpoint)}') model.discriminator.load_weights(str(latest_disc_checkpoint)) - - def save_model(step): if step % args.save_every == 0: print(f'Saving model on step {step} to {model_path}') @@ -127,8 +125,8 @@ def save_model(step): pad_range = (39, 49) time_range = (266, 276) data, features = preprocessing.read_csv_2d(pad_range=pad_range, time_range=time_range) - assert np.isclose(features[:,1].std(), 0), features[:,1].std() - features = features[:,:1].astype('float32') + assert np.isclose(features[:, 1].std(), 0), features[:, 1].std() + features = features[:, :1].astype('float32') if args.angles_to_radians: features *= np.pi / 180 @@ -139,7 +137,8 @@ def save_model(step): 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') - unscale = lambda x: 10 ** x - 1 + def unscale(x): + return 10 ** x - 1 def get_images(return_raw_data=False, calc_chi2=False, gen_more=None, sample=(X_test, Y_test), batch_size=128): X, Y = sample @@ -158,11 +157,11 @@ def get_images(return_raw_data=False, calc_chi2=False, gen_more=None, sample=(X_ gen = unscale(gen_scaled) gen[gen < 0] = 0 gen1 = np.where(gen < 1., 0, gen) - images = make_metric_plots(real, gen, features={'angle' : (X, gen_features)}, calc_chi2=calc_chi2) + images = make_metric_plots(real, gen, features={'angle': (X, gen_features)}, calc_chi2=calc_chi2) if calc_chi2: images, chi2 = images - images1 = make_metric_plots(real, gen1, features={'angle' : (X, gen_features)}) + images1 = make_metric_plots(real, gen1, features={'angle': (X, gen_features)}) img_amplitude = make_histograms(Y_test.flatten(), gen_scaled.flatten(), 'log10(amplitude + 1)', logy=True) @@ -176,7 +175,6 @@ def get_images(return_raw_data=False, calc_chi2=False, gen_more=None, sample=(X_ return result - def write_hist_summary(step): if step % args.save_every == 0: images, images1, img_amplitude, chi2 = get_images(calc_chi2=True) @@ -190,7 +188,6 @@ def write_hist_summary(step): tf.summary.image("{} (amp > 1)".format(k), img, step) tf.summary.image("log10(amplitude + 1)", img_amplitude, step) - def schedule_lr(step): model.disc_opt.lr.assign(model.disc_opt.lr * args.lr_schedule_rate) model.gen_opt.lr.assign(model.gen_opt.lr * args.lr_schedule_rate) @@ -203,7 +200,8 @@ def schedule_lr(step): assert not prediction_path.exists(), "Prediction path already exists" prediction_path.mkdir() - array_to_img = lambda arr: PIL.Image.fromarray(arr.reshape(arr.shape[1:])) + def array_to_img(arr): + return PIL.Image.fromarray(arr.reshape(arr.shape[1:])) for part in ['train', 'test']: path = prediction_path / part @@ -214,7 +212,7 @@ def schedule_lr(step): ) = get_images( calc_chi2=True, return_raw_data=True, gen_more=10, sample=( - (X_train, Y_train) if part=='train' + (X_train, Y_train) if part == 'train' else (X_test, Y_test) ) ) diff --git a/legacy_code/test_script_data_v3.py b/legacy_code/test_script_data_v3.py index 61016c4..83dd098 100644 --- a/legacy_code/test_script_data_v3.py +++ b/legacy_code/test_script_data_v3.py @@ -1,4 +1,5 @@ -import os, sys +import os +import sys import re from pathlib import Path import argparse @@ -37,10 +38,8 @@ def main(): parser.add_argument('--feature_noise_power', type=float, default=None) parser.add_argument('--feature_noise_decay', type=float, default=None) - args = parser.parse_args() - assert ( (args.feature_noise_power is None) == (args.feature_noise_decay is None) @@ -50,7 +49,7 @@ def main(): print("----" * 10) print("Arguments:") for k, v in vars(args).items(): - print(f" {k} : {v}") + print(f" {k}: {v}") print("----" * 10) print("") @@ -78,7 +77,8 @@ def main(): f.write('\n'.join(raw_args)) for fname in fnames: with open(fname, 'r') as f_in: - if len(raw_args) > 0: f.write('\n') + if len(raw_args) > 0: + f.write('\n') f.write(f_in.read()) model = BaselineModel_6x15(kernel_init=args.kernel_init, lr=args.lr, @@ -91,7 +91,7 @@ def main(): if args.prediction_only: def epoch_from_name(name): - epoch, = re.findall('\d+', name) + epoch, = re.findall(r'\d+', name) return int(epoch) gen_checkpoints = model_path.glob("generator_*.h5") @@ -114,8 +114,6 @@ def epoch_from_name(name): print(f'Loading discriminator weights from {str(latest_disc_checkpoint)}') model.discriminator.load_weights(str(latest_disc_checkpoint)) - - def save_model(step): if step % args.save_every == 0: print(f'Saving model on step {step} to {model_path}') @@ -135,7 +133,8 @@ def save_model(step): 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') - unscale = lambda x: 10 ** x - 1 + def unscale(x): + return 10 ** x - 1 def get_images(return_raw_data=False, calc_chi2=False, gen_more=None, sample=(X_test, Y_test), batch_size=128): X, Y = sample @@ -159,10 +158,10 @@ def get_images(return_raw_data=False, calc_chi2=False, gen_more=None, sample=(X_ gen1 = np.where(gen < 1., 0, gen) features = { - 'crossing_angle' : (X[:, 0], gen_features[:,0]), - 'dip_angle' : (X[:, 1], gen_features[:,1]), - 'drift_length' : (X[:, 2], gen_features[:,2]), - 'time_bin_fraction' : (X[:, 2] % 1, gen_features[:,2] % 1), + 'crossing_angle': (X[:, 0], gen_features[:, 0]), + 'dip_angle': (X[:, 1], gen_features[:, 1]), + 'drift_length': (X[:, 2], gen_features[:, 2]), + 'time_bin_fraction': (X[:, 2] % 1, gen_features[:, 2] % 1), } images = make_metric_plots(real, gen, features=features, calc_chi2=calc_chi2) @@ -183,7 +182,6 @@ def get_images(return_raw_data=False, calc_chi2=False, gen_more=None, sample=(X_ return result - def write_hist_summary(step): if step % args.save_every == 0: images, images1, img_amplitude, chi2 = get_images(calc_chi2=True) @@ -197,7 +195,6 @@ def write_hist_summary(step): tf.summary.image("{} (amp > 1)".format(k), img, step) tf.summary.image("log10(amplitude + 1)", img_amplitude, step) - def schedule_lr(step): model.disc_opt.lr.assign(model.disc_opt.lr * args.lr_schedule_rate) model.gen_opt.lr.assign(model.gen_opt.lr * args.lr_schedule_rate) @@ -206,12 +203,13 @@ def schedule_lr(step): tf.summary.scalar("generator learning rate", model.gen_opt.lr, step) if args.prediction_only: + def array_to_img(arr): + return PIL.Image.fromarray(arr.reshape(arr.shape[1:])) + prediction_path = model_path / f"prediction_{epoch_from_name(latest_gen_checkpoint.stem):05d}" assert not prediction_path.exists(), "Prediction path already exists" prediction_path.mkdir() - array_to_img = lambda arr: PIL.Image.fromarray(arr.reshape(arr.shape[1:])) - for part in ['train', 'test']: path = prediction_path / part path.mkdir() @@ -220,10 +218,7 @@ def schedule_lr(step): gen_dataset, chi2 ) = get_images( calc_chi2=True, return_raw_data=True, gen_more=10, - sample=( - (X_train, Y_train) if part=='train' - else (X_test, Y_test) - ) + sample=((X_train, Y_train) if part == 'train' else (X_test, Y_test)) ) for k, img in images.items(): array_to_img(img).save(str(path / f"{k}.png")) diff --git a/legacy_code/test_script_data_v4.py b/legacy_code/test_script_data_v4.py index 22b38e7..6b0ec29 100644 --- a/legacy_code/test_script_data_v4.py +++ b/legacy_code/test_script_data_v4.py @@ -1,4 +1,4 @@ -import os, sys +import sys import re from pathlib import Path import argparse @@ -14,6 +14,7 @@ from metrics import make_metric_plots, make_histograms import cuda_gpu_config + def make_parser(): parser = argparse.ArgumentParser(fromfile_prefix_chars='@') parser.add_argument('--checkpoint_name', type=str, required=True) @@ -44,7 +45,7 @@ def print_args(args): print("----" * 10) print("Arguments:") for k, v in vars(args).items(): - print(f" {k} : {v}") + print(f" {k}: {v}") print("----" * 10) print("") @@ -75,7 +76,7 @@ def write_args(model_path, fname='arguments.txt'): def epoch_from_name(name): - epoch, = re.findall('\d+', name) + epoch, = re.findall(r'\d+', name) return int(epoch) @@ -99,7 +100,7 @@ def load_weights(model, model_path): model.generator.load_weights(str(latest_gen_checkpoint)) print(f'Loading discriminator weights from {str(latest_disc_checkpoint)}') model.discriminator.load_weights(str(latest_disc_checkpoint)) - + return latest_gen_checkpoint, latest_disc_checkpoint @@ -130,11 +131,11 @@ def get_images(model, gen1 = np.where(gen < 1., 0, gen) features = { - 'crossing_angle' : (X[:, 0], gen_features[:,0]), - 'dip_angle' : (X[:, 1], gen_features[:,1]), - 'drift_length' : (X[:, 2], gen_features[:,2]), - 'time_bin_fraction' : (X[:, 2] % 1, gen_features[:,2] % 1), - 'pad_coord_fraction' : (X[:, 3] % 1, gen_features[:,3] % 1) + 'crossing_angle': (X[:, 0], gen_features[:, 0]), + 'dip_angle': (X[:, 1], gen_features[:, 1]), + 'drift_length': (X[:, 2], gen_features[:, 2]), + 'time_bin_fraction': (X[:, 2] % 1, gen_features[:, 2] % 1), + 'pad_coord_fraction': (X[:, 3] % 1, gen_features[:, 3] % 1) } images = make_metric_plots(real, gen, features=features, calc_chi2=calc_chi2) @@ -161,7 +162,7 @@ def __init__(self, model, path, save_period): self.model = model self.path = path self.save_period = save_period - + def __call__(self, step): if step % self.save_period == 0: print(f'Saving model on step {step} to {self.path}') @@ -208,6 +209,9 @@ def __call__(self, step): def evaluate_model(model, path, sample, gen_sample_name=None): + def array_to_img(arr): + return PIL.Image.fromarray(arr.reshape(arr.shape[1:])) + path.mkdir() ( images, images1, img_amplitude, @@ -215,8 +219,6 @@ def evaluate_model(model, path, sample, gen_sample_name=None): ) = get_images(model, sample=sample, calc_chi2=True, return_raw_data=True, gen_more=10) - array_to_img = lambda arr: PIL.Image.fromarray(arr.reshape(arr.shape[1:])) - for k, img in images.items(): array_to_img(img).save(str(path / f"{k}.png")) for k, img in images1.items(): @@ -276,7 +278,6 @@ def main(): 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') - if args.prediction_only: prediction_path = model_path / f"prediction_{epoch_from_name(latest_gen_checkpoint.stem):05d}" assert not prediction_path.exists(), "Prediction path already exists" @@ -302,7 +303,6 @@ def features_noise(epoch): return current_power - save_model = SaveModelCallback( model=model, path=model_path, save_period=args.save_every ) diff --git a/metrics/__init__.py b/metrics/__init__.py index 6d9771d..4a1e8be 100644 --- a/metrics/__init__.py +++ b/metrics/__init__.py @@ -2,49 +2,50 @@ import numpy as np import matplotlib as mpl -mpl.use("Agg") import matplotlib.pyplot as plt import PIL - -from .plotting import _bootstrap_error from .gaussian_metrics import get_val_metric_v, _METRIC_NAMES from .trends import make_trend_plot +mpl.use("Agg") + def make_histograms(data_real, data_gen, title, figsize=(8, 8), n_bins=100, logy=False, pdffile=None): - l = min(data_real.min(), data_gen.min()) - r = max(data_real.max(), data_gen.max()) - bins = np.linspace(l, r, n_bins + 1) - + left = min(data_real.min(), data_gen.min()) + right = max(data_real.max(), data_gen.max()) + bins = np.linspace(left, right, n_bins + 1) + fig = plt.figure(figsize=figsize) plt.hist(data_real, bins=bins, density=True, label='real') - plt.hist(data_gen , bins=bins, density=True, label='generated', alpha=0.7) + plt.hist(data_gen, bins=bins, density=True, label='generated', alpha=0.7) if logy: plt.yscale('log') plt.legend() plt.title(title) - + buf = io.BytesIO() fig.savefig(buf, format='png') - if pdffile is not None: fig.savefig(pdffile, format='pdf') + if pdffile is not None: + fig.savefig(pdffile, format='pdf') plt.close(fig) buf.seek(0) - + img = PIL.Image.open(buf) return np.array(img.getdata(), dtype=np.uint8).reshape(1, img.size[1], img.size[0], -1) def make_metric_plots(images_real, images_gen, features=None, calc_chi2=False, make_pdfs=False): plots = {} - if make_pdfs: pdf_plots = {} + if make_pdfs: + pdf_plots = {} if calc_chi2: chi2 = 0 try: metric_real = get_val_metric_v(images_real) - metric_gen = get_val_metric_v(images_gen ) - + metric_gen = get_val_metric_v(images_gen) + for name, real, gen in zip(_METRIC_NAMES, metric_real.T, metric_gen.T): pdffile = None if make_pdfs: @@ -52,7 +53,6 @@ def make_metric_plots(images_real, images_gen, features=None, calc_chi2=False, m pdf_plots[name] = pdffile plots[name] = make_histograms(real, gen, name, pdffile=pdffile) - if features is not None: for feature_name, (feature_real, feature_gen) in features.items(): for metric_name, real, gen in zip(_METRIC_NAMES, metric_real.T, metric_gen.T): @@ -74,7 +74,7 @@ def make_metric_plots(images_real, images_gen, features=None, calc_chi2=False, m except AssertionError as e: print(f"WARNING! Assertion error ({e})") - result = {'plots' : plots} + result = {'plots': plots} if calc_chi2: result['chi2'] = chi2 if make_pdfs: @@ -115,11 +115,11 @@ def make_images_for_model(model, gen1 = np.where(gen < 1., 0, gen) features = { - 'crossing_angle' : (X[:, 0], gen_features[:,0]), - 'dip_angle' : (X[:, 1], gen_features[:,1]), - 'drift_length' : (X[:, 2], gen_features[:,2]), - 'time_bin_fraction' : (X[:, 2] % 1, gen_features[:,2] % 1), - 'pad_coord_fraction' : (X[:, 3] % 1, gen_features[:,3] % 1) + 'crossing_angle': (X[:, 0], gen_features[:, 0]), + 'dip_angle': (X[:, 1], gen_features[:, 1]), + 'drift_length': (X[:, 2], gen_features[:, 2]), + 'time_bin_fraction': (X[:, 2] % 1, gen_features[:, 2] % 1), + 'pad_coord_fraction': (X[:, 3] % 1, gen_features[:, 3] % 1) } metric_plot_results = make_metric_plots(real, gen, features=features, @@ -165,6 +165,9 @@ def make_images_for_model(model, def evaluate_model(model, path, sample, gen_sample_name=None): + def array_to_img(arr): + return PIL.Image.fromarray(arr.reshape(arr.shape[1:])) + path.mkdir() pdf_outputs = [] ( @@ -174,8 +177,6 @@ def evaluate_model(model, path, sample, gen_sample_name=None): calc_chi2=True, return_raw_data=True, gen_more=10, pdf_outputs=pdf_outputs) images_pdf, images1_pdf, img_amplitude_pdf = pdf_outputs - array_to_img = lambda arr: PIL.Image.fromarray(arr.reshape(arr.shape[1:])) - for k, img in images.items(): array_to_img(img).save(str(path / f"{k}.png")) for k, img in images1.items(): @@ -234,7 +235,8 @@ def plot_individual_images(real, gen, n=10, pdffile=None): buf = io.BytesIO() fig.savefig(buf, format='png') - if pdffile is not None: fig.savefig(pdffile, format='pdf') + if pdffile is not None: + fig.savefig(pdffile, format='pdf') plt.close(fig) buf.seek(0) @@ -257,9 +259,10 @@ def plot_images_mask(real, gen, pdffile=None): buf = io.BytesIO() fig.savefig(buf, format='png') - if pdffile is not None: fig.savefig(pdffile, format='pdf') + if pdffile is not None: + fig.savefig(pdffile, format='pdf') plt.close(fig) buf.seek(0) img = PIL.Image.open(buf) - return np.array(img.getdata(), dtype=np.uint8).reshape(1, img.size[1], img.size[0], -1) \ No newline at end of file + return np.array(img.getdata(), dtype=np.uint8).reshape(1, img.size[1], img.size[0], -1) diff --git a/metrics/gaussian_metrics.py b/metrics/gaussian_metrics.py index af490d0..718691f 100644 --- a/metrics/gaussian_metrics.py +++ b/metrics/gaussian_metrics.py @@ -8,12 +8,12 @@ def _gaussian_fit(img): img_n = img / img.sum() mu = np.fromfunction( - lambda i, j: (img_n[np.newaxis,...] * np.stack([i, j])).sum(axis=(1, 2)), + lambda i, j: (img_n[np.newaxis, ...] * np.stack([i, j])).sum(axis=(1, 2)), shape=img.shape ) cov = np.fromfunction( lambda i, j: ( - (img_n[np.newaxis,...] * np.stack([i * i, j * i, i * j, j * j])).sum(axis=(1, 2)) + (img_n[np.newaxis, ...] * np.stack([i * i, j * i, i * j, j * j])).sum(axis=(1, 2)) ) - np.stack([mu[0]**2, mu[0]*mu[1], mu[0]*mu[1], mu[1]**2]), shape=img.shape ).reshape(2, 2) @@ -48,15 +48,15 @@ def get_val_metric_v(imgs): assert (imgs > 0).any(axis=(1, 2)).all(), 'get_val_metric_v: some images are empty' imgs_n = imgs / imgs.sum(axis=(1, 2), keepdims=True) mu = np.fromfunction( - lambda i, j: (imgs_n[:,np.newaxis,...] * np.stack([i, j])[np.newaxis,...]).sum(axis=(2, 3)), + lambda i, j: (imgs_n[:, np.newaxis, ...] * np.stack([i, j])[np.newaxis, ...]).sum(axis=(2, 3)), shape=imgs.shape[1:] ) cov = np.fromfunction( lambda i, j: ( - (imgs_n[:,np.newaxis,...] * np.stack([i * i, j * j, i * j])[np.newaxis,...]).sum(axis=(2, 3)) - ) - np.stack([mu[:,0]**2, mu[:,1]**2, mu[:,0] * mu[:,1]]).T, + (imgs_n[:, np.newaxis, ...] * np.stack([i * i, j * j, i * j])[np.newaxis, ...]).sum(axis=(2, 3)) + ) - np.stack([mu[:, 0]**2, mu[:, 1]**2, mu[:, 0] * mu[:, 1]]).T, shape=imgs.shape[1:] ) - return np.concatenate([mu, cov, imgs.sum(axis=(1, 2))[:,np.newaxis]], axis=1) \ No newline at end of file + return np.concatenate([mu, cov, imgs.sum(axis=(1, 2))[:, np.newaxis]], axis=1) diff --git a/metrics/plotting.py b/metrics/plotting.py index 26a114d..91f5761 100644 --- a/metrics/plotting.py +++ b/metrics/plotting.py @@ -14,13 +14,13 @@ def _get_stats(arr): class Obj: pass result = Obj() - + result.mean = arr.mean() result.width = arr.std() - + result.mean_err = result.width / (len(arr) - 1)**0.5 result.width_err = _bootstrap_error(arr, np.std) - + return result @@ -31,28 +31,28 @@ def compare_two_dists(d_real, d_gen, label, tag=None, nbins=100): max(d_real.max(), d_gen.max()), nbins + 1 ) - + stats_real = _get_stats(d_real) - stats_gen = _get_stats(d_gen ) - + stats_gen = _get_stats(d_gen) + if tag: leg_entry = f'gen ({tag})' else: leg_entry = 'gen' - + plt.hist(d_real, bins=bins, density=True, label='real') plt.hist(d_gen, bins=bins, density=True, label=leg_entry, histtype='step', linewidth=2.) - + string = '\n'.join([ f"real: mean = {stats_real.mean :.4f} +/- {stats_real.mean_err :.4f}", - f"gen: mean = {stats_gen .mean :.4f} +/- {stats_gen .mean_err :.4f}", + f"gen: mean = {stats_gen.mean :.4f} +/- {stats_gen .mean_err :.4f}", f"real: std = {stats_real.width:.4f} +/- {stats_real.width_err:.4f}", - f"gen: std = {stats_gen .width:.4f} +/- {stats_gen .width_err:.4f}", + f"gen: std = {stats_gen.width:.4f} +/- {stats_gen .width_err:.4f}", ]) default_family = rcParams['font.family'] rcParams['font.family'] = 'monospace' ax.add_artist(AnchoredText(string, loc=2)) rcParams['font.family'] = default_family - + plt.xlabel(label) plt.legend() diff --git a/metrics/trends.py b/metrics/trends.py index 2ea7d5c..42dab58 100644 --- a/metrics/trends.py +++ b/metrics/trends.py @@ -18,8 +18,8 @@ def calc_trend(x, y, do_plot=True, bins=100, window_size=20, **kwargs): bins = np.linspace(np.min(x), np.max(x), bins + 1) sel = (x >= bins[0]) x, y = x[sel], y[sel] - cats = (x[:,np.newaxis] < bins[np.newaxis,1:]).argmax(axis=1) - + cats = (x[:, np.newaxis] < bins[np.newaxis, 1:]).argmax(axis=1) + def stats(arr): return ( arr.mean(), @@ -27,7 +27,7 @@ def stats(arr): arr.std(), _bootstrap_error(arr, np.std) ) - + mean, mean_err, std, std_err, bin_centers = np.array([ stats( y[(cats >= left) & (cats < right)] @@ -37,12 +37,11 @@ def stats(arr): ) ]).T - if do_plot: mean_p_std_err = (mean_err**2 + std_err**2)**0.5 plt.fill_between(bin_centers, mean - mean_err, mean + mean_err, **kwargs) kwargs['alpha'] *= 0.5 - kwargs = {k : v for k, v in kwargs.items() if k != 'label'} + kwargs = {k: v for k, v in kwargs.items() if k != 'label'} plt.fill_between(bin_centers, mean - std - mean_p_std_err, mean - std + mean_p_std_err, **kwargs) plt.fill_between(bin_centers, mean + std - mean_p_std_err, mean + std + mean_p_std_err, **kwargs) kwargs['alpha'] *= 0.25 @@ -71,10 +70,11 @@ def make_trend_plot(feature_real, real, feature_gen, gen, name, calc_chi2=False, buf = io.BytesIO() fig.savefig(buf, format='png') - if pdffile is not None: fig.savefig(pdffile, format='pdf') + if pdffile is not None: + fig.savefig(pdffile, format='pdf') plt.close(fig) buf.seek(0) - + img = PIL.Image.open(buf) img_data = np.array(img.getdata(), dtype=np.uint8).reshape(1, img.size[1], img.size[0], -1) @@ -107,5 +107,5 @@ def make_trend_plot(feature_real, real, feature_gen, gen, name, calc_chi2=False, ) return img_data, chi2 - - return img_data \ No newline at end of file + + return img_data diff --git a/model_export/dump_graph.py b/model_export/dump_graph.py index 23e52e2..8852cfe 100644 --- a/model_export/dump_graph.py +++ b/model_export/dump_graph.py @@ -7,10 +7,11 @@ from tensorflow.core.framework import attr_value_pb2 from tensorflow.python.tools import optimize_for_inference_lib -from . import tf2xla_pb2 +import tf2xla_pb2 + def model_to_graph(model, preprocess, postprocess, input_signature, output_file, test_input=None, - hack_upsampling=False, batch_sizes=(1, 10, 100, 1000, 10000), perf_iterations=5): + hack_upsampling=False, batch_sizes=(1, 10, 100, 1000, 10000), perf_iterations=5): tf.keras.backend.set_learning_phase(0) @tf.function(input_signature=input_signature) @@ -31,7 +32,6 @@ def to_save(x): op._set_attr('align_corners', attr_value_pb2.AttrValue(b=True)) op._set_attr('half_pixel_centers', attr_value_pb2.AttrValue(b=False)) - output_file = Path(output_file) path = str(output_file.parent) filename = output_file.name @@ -73,13 +73,13 @@ def to_save(x): for batch_size in batch_sizes[::-1]: timings = [] iterations = perf_iterations * max(1, 100 // batch_size) - for i in range(perf_iterations): + for i in range(iterations): batched_input = tf.random.normal( shape=(batch_size, len(test_input)), dtype='float32' ) t0 = perf_counter() - result = to_save(batched_input).numpy() + to_save(batched_input).numpy() t1 = perf_counter() timings.append((t1 - t0) * 1000. / batch_size) @@ -87,5 +87,3 @@ def to_save(x): mean = timings.mean() err = timings.std() / (len(timings) - 1)**0.5 print(f'With batch size = {batch_size}, duration per 1 generation is: {mean} +\\- {err} ms') - - diff --git a/models/callbacks.py b/models/callbacks.py index e447d57..66b5b3b 100644 --- a/models/callbacks.py +++ b/models/callbacks.py @@ -2,6 +2,7 @@ from metrics import make_images_for_model + class SaveModelCallback: def __init__(self, model, path, save_period): self.model = model @@ -60,4 +61,4 @@ def get_scheduler(lr, lr_decay): def schedule_lr(step): return lr * lr_decay**step - return schedule_lr \ No newline at end of file + return schedule_lr diff --git a/models/model_v4.py b/models/model_v4.py index ded6544..788850d 100644 --- a/models/model_v4.py +++ b/models/model_v4.py @@ -5,6 +5,7 @@ from . import scalers, nn + @tf.function(experimental_relax_shapes=True) def preprocess_features(features): # features: @@ -12,14 +13,16 @@ def preprocess_features(features): # dip_angle [-60, 60] # drift_length [35, 290] # pad_coordinate [40-something, 40-something] - bin_fractions = features[:,-2:] % 1 + bin_fractions = features[:, -2:] % 1 features = ( - features[:,:3] - tf.constant([[0., 0., 162.5]]) + features[:, :3] - tf.constant([[0., 0., 162.5]]) ) / tf.constant([[20., 60., 127.5]]) return tf.concat([features, bin_fractions], axis=-1) + _f = preprocess_features + def disc_loss(d_real, d_fake): return tf.reduce_mean(d_fake - d_real) @@ -31,11 +34,12 @@ def gen_loss(d_real, d_fake): def disc_loss_cramer(d_real, d_fake, d_fake_2): return -tf.reduce_mean( tf.norm(d_real - d_fake, axis=-1) + - tf.norm(d_fake_2, axis=-1) - + tf.norm(d_fake_2, axis=-1) - tf.norm(d_fake - d_fake_2, axis=-1) - tf.norm(d_real, axis=-1) ) + def gen_loss_cramer(d_real, d_fake, d_fake_2): return -disc_loss_cramer(d_real, d_fake, d_fake_2) @@ -51,6 +55,7 @@ def disc_loss_js(d_real, d_fake): logloss(-d_fake) ) / (len(d_real) + len(d_fake)) + def gen_loss_js(d_real, d_fake): return tf.reduce_mean( logloss(d_fake) @@ -126,7 +131,6 @@ def _load_weights(self, checkpoint, gen_or_disc): ) network.optimizer.set_weights(opt_weight_values) - @tf.function def make_fake(self, features): size = tf.shape(features)[0] @@ -136,7 +140,7 @@ def make_fake(self, features): ) def gradient_penalty(self, features, real, fake): - alpha = tf.random.uniform(shape=[len(real),] + [1] * (len(real.shape) - 1)) + alpha = tf.random.uniform(shape=[len(real), ] + [1] * (len(real.shape) - 1)) interpolates = alpha * real + (1 - alpha) * fake with tf.GradientTape() as t: t.watch(interpolates) diff --git a/models/nn.py b/models/nn.py index f063306..c55c391 100644 --- a/models/nn.py +++ b/models/nn.py @@ -1,5 +1,4 @@ import tensorflow as tf -import numpy as np custom_objects = {} @@ -148,7 +147,7 @@ def vector_img_connect_block(vector_shape, img_shape, block, block_input = input_img if len(img_shape) == 2: block_input = tf.keras.layers.Reshape(img_shape + (1,))(block_input) - if not vector_bypass: + if not vector_bypass: reshaped_vec = tf.tile( tf.keras.layers.Reshape((1, 1) + vector_shape)(input_vec), (1, *img_shape[:2], 1) @@ -214,4 +213,4 @@ def build_architecture(block_descriptions, name=None, custom_objects_code=None): ) if name: args['name'] = name - return tf.keras.Model(**args) \ No newline at end of file + return tf.keras.Model(**args) diff --git a/models/scalers.py b/models/scalers.py index 83969e5..98cca0f 100644 --- a/models/scalers.py +++ b/models/scalers.py @@ -6,7 +6,7 @@ class Identity: def scale(self, x): return x - + def unscale(self, x): return x @@ -25,8 +25,8 @@ def __init__(self, shape=(8, 16)): def scale(self, x): result = gaussian_fit(x) - result[:,-1] = np.log1p(result[:,-1]) - result[:,4] /= (result[:,2] * result[:,3]) + result[:, -1] = np.log1p(result[:, -1]) + result[:, 4] /= (result[:, 2] * result[:, 3]) return result def unscale(self, x): @@ -41,7 +41,7 @@ def unscale(self, x): cov = np.stack([ np.stack([D00, D01], axis=1), np.stack([D01, D11], axis=1) - ], axis=2) # N x 2 x 2 + ], axis=2) # N x 2 x 2 invcov = np.linalg.inv(cov) mu = np.stack([m0, m1], axis=1) @@ -49,14 +49,12 @@ def unscale(self, x): xx1 = np.arange(self.shape[1]) xx0, xx1 = np.meshgrid(xx0, xx1, indexing='ij') xx = np.stack([xx0, xx1], axis=2) - residuals = xx[None,...] - mu[:,None,None,:] # N x H x W x 2 + residuals = xx[None, ...] - mu[:, None, None, :] # N x H x W x 2 - result = np.exp(-0.5 * - np.einsum('ijkl,ilm,ijkm->ijk', residuals, invcov, residuals) - ) + result = np.exp(-0.5 * np.einsum('ijkl,ilm,ijkm->ijk', residuals, invcov, residuals)) result /= result.sum(axis=(1, 2), keepdims=True) - result *= A[:,None,None] + result *= A[:, None, None] return result diff --git a/models/training.py b/models/training.py index afcf3e6..59c59c5 100644 --- a/models/training.py +++ b/models/training.py @@ -14,7 +14,6 @@ def train(data_train, data_val, train_step_fn, loss_eval_fn, num_epochs, batch_s print("Working on epoch #{}".format(i_epoch), flush=True) tf.keras.backend.set_learning_phase(1) # training - shuffle_ids = np.random.permutation(len(data_train)) losses_train = {} @@ -40,7 +39,7 @@ def train(data_train, data_val, train_step_fn, loss_eval_fn, num_epochs, batch_s losses_train_batch = train_step_fn(feature_batch, batch) for k, l in losses_train_batch.items(): losses_train[k] = losses_train.get(k, 0) + l.numpy() * len(batch) - losses_train = {k : l / len(data_train) for k, l in losses_train.items()} + losses_train = {k: l / len(data_train) for k, l in losses_train.items()} tf.keras.backend.set_learning_phase(0) # testing @@ -49,13 +48,13 @@ def train(data_train, data_val, train_step_fn, loss_eval_fn, num_epochs, batch_s batch = data_val[i_sample:i_sample + batch_size] if features_train is None: - losses_val_batch = {k : l.numpy() for k, l in loss_eval_fn(batch).items()} + losses_val_batch = {k: l.numpy() for k, l in loss_eval_fn(batch).items()} else: feature_batch = features_val[i_sample:i_sample + batch_size] - losses_val_batch = {k : l.numpy() for k, l in loss_eval_fn(feature_batch, batch).items()} + losses_val_batch = {k: l.numpy() for k, l in loss_eval_fn(feature_batch, batch).items()} for k, l in losses_val_batch.items(): losses_val[k] = losses_val.get(k, 0) + l * len(batch) - losses_val = {k : l / len(data_val) for k, l in losses_val.items()} + losses_val = {k: l / len(data_val) for k, l in losses_val.items()} for f in callbacks: f(i_epoch) @@ -64,7 +63,7 @@ def train(data_train, data_val, train_step_fn, loss_eval_fn, num_epochs, batch_s with train_writer.as_default(): for k, l in losses_train.items(): tf.summary.scalar(k, l, i_epoch) - + if val_writer is not None: with val_writer.as_default(): for k, l in losses_val.items(): diff --git a/models/utils.py b/models/utils.py index dd5a41b..98743b0 100644 --- a/models/utils.py +++ b/models/utils.py @@ -2,7 +2,7 @@ def epoch_from_name(name): - epoch, = re.findall('\d+', name) + epoch, = re.findall(r'\d+', name) return int(epoch) @@ -33,4 +33,4 @@ def load_weights(model, model_path, epoch=None): model.load_generator(gen_checkpoint) model.load_discriminator(disc_checkpoint) - return epoch \ No newline at end of file + return epoch diff --git a/run_model_v4.py b/run_model_v4.py index c434129..73ed5da 100644 --- a/run_model_v4.py +++ b/run_model_v4.py @@ -1,12 +1,9 @@ -import os, sys from pathlib import Path import shutil import argparse -import numpy as np from sklearn.model_selection import train_test_split import tensorflow as tf -import PIL import yaml from data import preprocessing @@ -17,6 +14,7 @@ from metrics import evaluate_model import cuda_gpu_config + def make_parser(): parser = argparse.ArgumentParser(fromfile_prefix_chars='@') parser.add_argument('--config', type=str, required=False) @@ -28,13 +26,13 @@ def make_parser(): def print_args(args): - print("") + print() print("----" * 10) print("Arguments:") for k, v in vars(args).items(): print(f" {k} : {v}") print("----" * 10) - print("") + print() def parse_args(): @@ -52,10 +50,14 @@ def load_config(file): (config['feature_noise_decay'] is None) ), 'Noise power and decay must be both provided' - if 'lr_disc' not in config: config['lr_disc'] = config['lr'] - if 'lr_gen' not in config: config['lr_gen' ] = config['lr'] - if 'lr_schedule_rate_disc' not in config: config['lr_schedule_rate_disc'] = config['lr_schedule_rate'] - if 'lr_schedule_rate_gen' not in config: config['lr_schedule_rate_gen' ] = config['lr_schedule_rate'] + if 'lr_disc' not in config: + config['lr_disc'] = config['lr'] + if 'lr_gen' not in config: + config['lr_gen'] = config['lr'] + if 'lr_schedule_rate_disc' not in config: + config['lr_schedule_rate_disc'] = config['lr_schedule_rate'] + if 'lr_schedule_rate_gen' not in config: + config['lr_schedule_rate_gen'] = config['lr_schedule_rate'] return config @@ -103,7 +105,6 @@ def main(): 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') - if args.prediction_only: epoch = latest_epoch(model_path) prediction_path = model_path / f"prediction_{epoch:05d}" @@ -130,7 +131,6 @@ def features_noise(epoch): return current_power - save_model = SaveModelCallback( model=model, path=model_path, save_period=config['save_every'] ) From 7a11afe84192c23096d8cdbb30684abdaf233ac1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A1=D1=83=D1=85=D0=BE=D1=80=D0=BE=D1=81=D0=BE=D0=B2=20?= =?UTF-8?q?=D0=90=D0=BB=D0=B5=D0=BA=D1=81=D0=B5=D0=B9?= Date: Wed, 6 Oct 2021 22:48:51 +0300 Subject: [PATCH 3/6] add black --- .github/workflows/push.yml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 62db9e7..9a65522 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -1,4 +1,4 @@ -name: Flake8 +name: linter on: [push] jobs: flake8: @@ -13,3 +13,8 @@ jobs: run: python -m pip install flake8 - name: Run flake8 run: flake8 + black: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + - uses: psf/black@stable \ No newline at end of file From 137deaa180e9cf748f00242e38894ee1cbcac563 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A1=D1=83=D1=85=D0=BE=D1=80=D0=BE=D1=81=D0=BE=D0=B2=20?= =?UTF-8?q?=D0=90=D0=BB=D0=B5=D0=BA=D1=81=D0=B5=D0=B9?= Date: Wed, 6 Oct 2021 23:02:05 +0300 Subject: [PATCH 4/6] black stylefixes --- .flake8 | 1 + combine_images.py | 22 +- data/preprocessing.py | 23 +- dump_graph_model_v4.py | 19 +- legacy_code/models/baseline_10x10.py | 142 ++-- legacy_code/models/baseline_10x15.py | 89 ++- legacy_code/models/baseline_fc_v4_8x16.py | 194 +++-- legacy_code/models/baseline_v2_10x10.py | 211 +++--- legacy_code/models/baseline_v3_6x15.py | 229 +++--- legacy_code/models/baseline_v4_8x16.py | 244 +++---- legacy_code/test_script_data_v0.py | 20 +- legacy_code/test_script_data_v1.py | 21 +- legacy_code/test_script_data_v1_normed.py | 22 +- legacy_code/test_script_data_v2.py | 99 +-- legacy_code/test_script_data_v3.py | 97 +-- legacy_code/test_script_data_v4.py | 127 ++-- metrics/__init__.py | 62 +- metrics/gaussian_metrics.py | 21 +- metrics/plotting.py | 27 +- metrics/trends.py | 60 +- model_export/dump_graph.py | 42 +- model_export/tf2xla_pb2.py | 830 ++++++++++++++-------- models/callbacks.py | 13 +- models/model_v4.py | 65 +- models/nn.py | 74 +- models/scalers.py | 9 +- models/training.py | 32 +- models/utils.py | 6 +- pyproject.toml | 3 + run_model_v4.py | 50 +- 30 files changed, 1529 insertions(+), 1325 deletions(-) create mode 100644 pyproject.toml diff --git a/.flake8 b/.flake8 index 9fd30a6..34eda72 100644 --- a/.flake8 +++ b/.flake8 @@ -1,5 +1,6 @@ [flake8] max-line-length = 120 +extend-ignore = E203 exclude = model_export/tf2xla_pb2.py per-file-ignores = diff --git a/combine_images.py b/combine_images.py index 41d0dc1..66fddce 100644 --- a/combine_images.py +++ b/combine_images.py @@ -12,20 +12,20 @@ def main(): args = parser.parse_args() variables = [ - 'crossing_angle', - 'dip_angle', - 'drift_length', - 'pad_coord_fraction', - 'time_bin_fraction', + 'crossing_angle', + 'dip_angle', + 'drift_length', + 'pad_coord_fraction', + 'time_bin_fraction', ] stats = [ - 'Mean0', - 'Mean1', - 'Sigma0^2', - 'Sigma1^2', - 'Cov01', - 'Sum', + 'Mean0', + 'Mean1', + 'Sigma0^2', + 'Sigma1^2', + 'Cov01', + 'Sum', ] img_path = Path(args.path_to_images) diff --git a/data/preprocessing.py b/data/preprocessing.py index a29c2c3..ca75fb0 100644 --- a/data/preprocessing.py +++ b/data/preprocessing.py @@ -38,10 +38,7 @@ def raw_to_csv(fname_in=None, fname_out=None): with open(fname_in, 'r') as f: lines = f.readlines() - reader_main = Reader( - variables=['ipad', 'itime', 'amp'], - types=[int, int, float] - ) + reader_main = Reader(variables=['ipad', 'itime', 'amp'], types=[int, int, float]) data_sources = [lines] readers = [reader_main] @@ -50,19 +47,15 @@ def raw_to_csv(fname_in=None, fname_out=None): assert len(lines) % 2 == 0, 'raw_to_csv: Odd number of lines when expected even' if _VERSION == 'data_v2': - reader_features = Reader( - variables=["crossing_angle", "dip_angle"], - types=[float, float] - ) + reader_features = Reader(variables=["crossing_angle", "dip_angle"], types=[float, float]) elif _VERSION == 'data_v3': reader_features = Reader( - variables=["crossing_angle", "dip_angle", "drift_length"], - types=[float, float, float] + variables=["crossing_angle", "dip_angle", "drift_length"], types=[float, float, float] ) elif _VERSION == 'data_v4': reader_features = Reader( variables=["crossing_angle", "dip_angle", "drift_length", "pad_coordinate"], - types=[float, float, float, float] + types=[float, float, float, float], ) else: raise NotImplementedError @@ -96,10 +89,7 @@ def sel(df, col, limits): if 'pad_coordinate' in df.columns: df['ipad'] -= df['pad_coordinate'].astype(int) - selection = ( - sel(df, 'itime', time_range) & - sel(df, 'ipad', pad_range) - ) + selection = sel(df, 'itime', time_range) & sel(df, 'ipad', pad_range) if not selection.all(): msg = f"WARNING: current selection ignores {(~selection).sum() / len(selection) * 100}% of the data!" @@ -109,8 +99,7 @@ def sel(df, col, limits): g = df[selection].groupby('evtId') def convert_event(event): - result = np.zeros(dtype=float, shape=(pad_range[1] - pad_range[0], - time_range[1] - time_range[0])) + result = np.zeros(dtype=float, shape=(pad_range[1] - pad_range[0], time_range[1] - time_range[0])) indices = tuple(event[['ipad', 'itime']].values.T - np.array([[pad_range[0]], [time_range[0]]])) result[indices] = event.amp.values diff --git a/dump_graph_model_v4.py b/dump_graph_model_v4.py index 2f7483c..5fe045a 100644 --- a/dump_graph_model_v4.py +++ b/dump_graph_model_v4.py @@ -1,5 +1,4 @@ import argparse -import re from pathlib import Path import tensorflow as tf @@ -33,10 +32,6 @@ def main(): print("----" * 10) print("") - def epoch_from_name(name): - epoch, = re.findall(r'\d+', name) - return int(epoch) - model_path = Path('saved_models') / args.checkpoint_name full_model = Model_v4(load_config(model_path / 'config.yaml')) @@ -44,11 +39,14 @@ def epoch_from_name(name): model = full_model.generator if args.constant_seed is None: + def preprocess(x): size = tf.shape(x)[0] latent_input = tf.random.normal(shape=(size, args.latent_dim), dtype='float32') return tf.concat([preprocess_features(x), latent_input], axis=-1) + else: + def preprocess(x): size = tf.shape(x)[0] latent_input = tf.ones(shape=(size, args.latent_dim), dtype='float32') * args.constant_seed @@ -56,13 +54,16 @@ def preprocess(x): def postprocess(x): x = 10 ** x - 1 - return tf.where(x < 1., 0., x) + return tf.where(x < 1.0, 0.0, x) dump_graph.model_to_graph( - model, preprocess, postprocess, + model, + preprocess, + postprocess, input_signature=[tf.TensorSpec(shape=[None, 4], dtype=tf.float32)], - output_file=args.output_path, test_input=args.test_input, - hack_upsampling=not args.dont_hack_upsampling_op + output_file=args.output_path, + test_input=args.test_input, + hack_upsampling=not args.dont_hack_upsampling_op, ) diff --git a/legacy_code/models/baseline_10x10.py b/legacy_code/models/baseline_10x10.py index 0c64cf5..1d0d62d 100644 --- a/legacy_code/models/baseline_10x10.py +++ b/legacy_code/models/baseline_10x10.py @@ -2,66 +2,74 @@ def get_generator(activation, kernel_init, latent_dim): - generator = tf.keras.Sequential([ - tf.keras.layers.Dense(units=64, activation=activation, input_shape=(latent_dim,)), - - tf.keras.layers.Reshape((4, 4, 4)), - - 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=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=8, kernel_size=3, padding='valid', activation=activation, - kernel_initializer=kernel_init), # 10x10 - tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, - kernel_initializer=kernel_init), - - tf.keras.layers.Reshape((10, 10)), - ], name='generator') + generator = tf.keras.Sequential( + [ + tf.keras.layers.Dense(units=64, activation=activation, input_shape=(latent_dim,)), + tf.keras.layers.Reshape((4, 4, 4)), + 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=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=8, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init + ), # 10x10 + tf.keras.layers.Conv2D( + filters=1, + kernel_size=1, + padding='valid', + activation=tf.keras.activations.relu, + kernel_initializer=kernel_init, + ), + tf.keras.layers.Reshape((10, 10)), + ], + name='generator', + ) return generator def get_discriminator(activation, kernel_init, dropout_rate): - discriminator = tf.keras.Sequential([ - tf.keras.layers.Reshape((10, 10, 1), input_shape=(10, 10)), - - 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(), # 4x4 - - 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.MaxPool2D(), # 2x2 - - 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.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') + discriminator = tf.keras.Sequential( + [ + tf.keras.layers.Reshape((10, 10, 1), input_shape=(10, 10)), + 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(), # 4x4 + 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.MaxPool2D(), # 2x2 + 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.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 @@ -74,8 +82,16 @@ def gen_loss(d_real, d_fake): class BaselineModel10x10: - def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uniform', - dropout_rate=0.2, lr=1e-4, latent_dim=32, gp_lambda=10., num_disc_updates=3): + def __init__( + self, + activation=tf.keras.activations.relu, + kernel_init='glorot_uniform', + dropout_rate=0.2, + lr=1e-4, + latent_dim=32, + gp_lambda=10.0, + num_disc_updates=3, + ): self.disc_opt = tf.keras.optimizers.RMSprop(lr) self.gen_opt = tf.keras.optimizers.RMSprop(lr) self.latent_dim = latent_dim @@ -84,17 +100,13 @@ def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uni self.generator = get_generator(activation=activation, kernel_init=kernel_init, latent_dim=latent_dim) self.discriminator = get_discriminator( - activation=activation, - kernel_init=kernel_init, - dropout_rate=dropout_rate + activation=activation, kernel_init=kernel_init, dropout_rate=dropout_rate ) self.step_counter = tf.Variable(0, dtype='int32', trainable=False) def make_fake(self, size): - return self.generator( - tf.random.normal(shape=(size, self.latent_dim), dtype='float32') - ) + return self.generator(tf.random.normal(shape=(size, self.latent_dim), dtype='float32')) def gradient_penalty(self, real, fake): alpha = tf.random.uniform(shape=[len(real), 1, 1]) @@ -103,7 +115,7 @@ def gradient_penalty(self, real, fake): t.watch(interpolates) d_int = self.discriminator(interpolates) 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) + return tf.reduce_mean(tf.maximum(tf.norm(grads, axis=-1) - 1, 0) ** 2) @tf.function def calculate_losses(self, batch): diff --git a/legacy_code/models/baseline_10x15.py b/legacy_code/models/baseline_10x15.py index b2099af..95ffa7c 100644 --- a/legacy_code/models/baseline_10x15.py +++ b/legacy_code/models/baseline_10x15.py @@ -1,54 +1,47 @@ import tensorflow as tf + LATENT_DIM = 32 activation = tf.keras.activations.elu dropout_rate = 0.2 NUM_DISC_UPDATES = 2 -GP_LAMBDA = 10. - -generator = tf.keras.Sequential([ - tf.keras.layers.Dense(units=64, activation=activation, input_shape=(LATENT_DIM,)), - - tf.keras.layers.Dense(units=480, activation=activation, input_shape=(LATENT_DIM,)), - tf.keras.layers.Reshape((3, 4, 40)), - - tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation), - tf.keras.layers.UpSampling2D(), # 6x8 - - tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='same', activation=activation), - tf.keras.layers.UpSampling2D(), # 12x16 - - tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='same', activation=activation), - tf.keras.layers.Conv2D(filters=4, kernel_size=(3, 2), padding='valid', activation=activation), - - tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=activation), - - tf.keras.layers.Reshape((10, 15)), -], name='generator') - -discriminator = tf.keras.Sequential([ - tf.keras.layers.Reshape((10, 15, 1), input_shape=(10, 15)), - - tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='same', activation=activation), - tf.keras.layers.Dropout(dropout_rate), - - tf.keras.layers.MaxPool2D(padding='same'), # 5x8 - - tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation), - tf.keras.layers.Dropout(dropout_rate), - - tf.keras.layers.MaxPool2D(padding='same'), # 3x4 - - tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation), # 1x2 - tf.keras.layers.Dropout(dropout_rate), - - tf.keras.layers.Reshape((128,)), - - tf.keras.layers.Dense(units=128, activation=activation), - tf.keras.layers.Dropout(dropout_rate), - - tf.keras.layers.Dense(units=1, activation=activation), -], name='discriminator') +GP_LAMBDA = 10.0 + +generator = tf.keras.Sequential( + [ + tf.keras.layers.Dense(units=64, activation=activation, input_shape=(LATENT_DIM,)), + tf.keras.layers.Dense(units=480, activation=activation, input_shape=(LATENT_DIM,)), + tf.keras.layers.Reshape((3, 4, 40)), + tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation), + tf.keras.layers.UpSampling2D(), # 6x8 + tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='same', activation=activation), + tf.keras.layers.UpSampling2D(), # 12x16 + tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='same', activation=activation), + tf.keras.layers.Conv2D(filters=4, kernel_size=(3, 2), padding='valid', activation=activation), + tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=activation), + tf.keras.layers.Reshape((10, 15)), + ], + name='generator', +) + +discriminator = tf.keras.Sequential( + [ + tf.keras.layers.Reshape((10, 15, 1), input_shape=(10, 15)), + tf.keras.layers.Conv2D(filters=16, kernel_size=3, padding='same', activation=activation), + tf.keras.layers.Dropout(dropout_rate), + tf.keras.layers.MaxPool2D(padding='same'), # 5x8 + tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same', activation=activation), + tf.keras.layers.Dropout(dropout_rate), + tf.keras.layers.MaxPool2D(padding='same'), # 3x4 + tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation), # 1x2 + tf.keras.layers.Dropout(dropout_rate), + tf.keras.layers.Reshape((128,)), + tf.keras.layers.Dense(units=128, activation=activation), + tf.keras.layers.Dropout(dropout_rate), + tf.keras.layers.Dense(units=1, activation=activation), + ], + name='discriminator', +) disc_opt = tf.optimizers.RMSprop() @@ -56,9 +49,7 @@ def make_fake(size): - return generator( - tf.random.normal(shape=(size, LATENT_DIM), dtype='float32') - ) + return generator(tf.random.normal(shape=(size, LATENT_DIM), dtype='float32')) def disc_loss(d_real, d_fake): @@ -76,7 +67,7 @@ def gradient_penalty(real, fake): t.watch(interpolates) d_int = discriminator(interpolates) 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) + return tf.reduce_mean(tf.maximum(tf.norm(grads, axis=-1) - 1, 0) ** 2) @tf.function diff --git a/legacy_code/models/baseline_fc_v4_8x16.py b/legacy_code/models/baseline_fc_v4_8x16.py index d466e8e..e27a143 100644 --- a/legacy_code/models/baseline_fc_v4_8x16.py +++ b/legacy_code/models/baseline_fc_v4_8x16.py @@ -11,9 +11,7 @@ def preprocess_features(features): # drift_length [35, 290] # pad_coordinate [40-something, 40-something] bin_fractions = features[:, -2:] % 1 - features = ( - features[:, :3] - tf.constant([[0., 0., 162.5]]) - ) / tf.constant([[20., 60., 127.5]]) + features = (features[:, :3] - tf.constant([[0.0, 0.0, 162.5]])) / tf.constant([[20.0, 60.0, 127.5]]) return tf.concat([features, bin_fractions], axis=-1) @@ -21,61 +19,63 @@ def preprocess_features(features): def get_generator(activation, kernel_init, num_features, latent_dim): - generator = tf.keras.Sequential([ - tf.keras.layers.Dense(units=32, activation=activation, input_shape=(num_features + latent_dim,)), - tf.keras.layers.Dense(units=64, activation=activation), - tf.keras.layers.Dense(units=64, activation=activation), - tf.keras.layers.Dense(units=64, activation=activation), - tf.keras.layers.Dense(units=8*16, activation=activation), - tf.keras.layers.Reshape((8, 16)), - ], name='generator') + generator = tf.keras.Sequential( + [ + tf.keras.layers.Dense(units=32, activation=activation, input_shape=(num_features + latent_dim,)), + tf.keras.layers.Dense(units=64, activation=activation), + tf.keras.layers.Dense(units=64, activation=activation), + tf.keras.layers.Dense(units=64, activation=activation), + tf.keras.layers.Dense(units=8 * 16, activation=activation), + tf.keras.layers.Reshape((8, 16)), + ], + name='generator', + ) return generator -def get_discriminator(activation, kernel_init, dropout_rate, num_features, num_additional_layers, cramer=False, - features_to_tail=False): +def get_discriminator( + activation, kernel_init, dropout_rate, num_features, num_additional_layers, cramer=False, features_to_tail=False +): input_img = tf.keras.Input(shape=(8, 16)) features_input = tf.keras.Input(shape=(num_features,)) img = tf.reshape(input_img, (-1, 8, 16, 1)) if features_to_tail: - features_tiled = tf.tile( - tf.reshape(features_input, (-1, 1, 1, num_features)), - (1, 8, 16, 1) - ) - img = tf.concat( - [img, features_tiled], - axis=-1 - ) - - discriminator_tail = tf.keras.Sequential([ - 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='same', activation=activation, - kernel_initializer=kernel_init), - tf.keras.layers.Dropout(dropout_rate), - - tf.keras.layers.MaxPool2D(pool_size=(1, 2)), # 8x8 - - 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.MaxPool2D(), # 4x4 - - tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation, - kernel_initializer=kernel_init), # 2x2 - 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.Reshape((64,)) - ], name='discriminator_tail') + features_tiled = tf.tile(tf.reshape(features_input, (-1, 1, 1, num_features)), (1, 8, 16, 1)) + img = tf.concat([img, features_tiled], axis=-1) + + discriminator_tail = tf.keras.Sequential( + [ + 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='same', activation=activation, kernel_initializer=kernel_init + ), + tf.keras.layers.Dropout(dropout_rate), + tf.keras.layers.MaxPool2D(pool_size=(1, 2)), # 8x8 + 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.MaxPool2D(), # 4x4 + tf.keras.layers.Conv2D( + filters=64, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init + ), # 2x2 + 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.Reshape((64,)), + ], + name='discriminator_tail', + ) head_input = tf.keras.layers.Concatenate()([features_input, discriminator_tail(img)]) @@ -90,19 +90,14 @@ def get_discriminator(activation, kernel_init, dropout_rate, num_features, num_a ] discriminator_head = tf.keras.Sequential( - head_layers + [tf.keras.layers.Dense(units=1 if not cramer else 256, - activation=None)], - name='discriminator_head' + head_layers + [tf.keras.layers.Dense(units=1 if not cramer else 256, activation=None)], + name='discriminator_head', ) inputs = [features_input, input_img] outputs = discriminator_head(head_input) - discriminator = tf.keras.Model( - inputs=inputs, - outputs=outputs, - name='discriminator' - ) + discriminator = tf.keras.Model(inputs=inputs, outputs=outputs, name='discriminator') return discriminator @@ -117,10 +112,10 @@ def gen_loss(d_real, d_fake): def disc_loss_cramer(d_real, d_fake, d_fake_2): return -tf.reduce_mean( - tf.norm(d_real - d_fake, axis=-1) + - tf.norm(d_fake_2, axis=-1) - - tf.norm(d_fake - d_fake_2, axis=-1) - - tf.norm(d_real, axis=-1) + tf.norm(d_real - d_fake, axis=-1) + + tf.norm(d_fake_2, axis=-1) + - tf.norm(d_fake - d_fake_2, axis=-1) + - tf.norm(d_real, axis=-1) ) @@ -129,10 +124,21 @@ def gen_loss_cramer(d_real, d_fake, d_fake_2): class BaselineModel_8x16: - def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uniform', - dropout_rate=0.02, lr=1e-4, latent_dim=32, gp_lambda=10., num_disc_updates=8, - gpdata_lambda=0., num_additional_layers=0, cramer=False, - features_to_tail=True, stochastic_stepping=True): + def __init__( + self, + activation=tf.keras.activations.relu, + kernel_init='glorot_uniform', + dropout_rate=0.02, + lr=1e-4, + latent_dim=32, + gp_lambda=10.0, + num_disc_updates=8, + gpdata_lambda=0.0, + num_additional_layers=0, + cramer=False, + features_to_tail=True, + stochastic_stepping=True, + ): self.disc_opt = tf.keras.optimizers.RMSprop(lr) self.gen_opt = tf.keras.optimizers.RMSprop(lr) self.latent_dim = latent_dim @@ -147,17 +153,22 @@ def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uni activation=activation, kernel_init=kernel_init, latent_dim=latent_dim, num_features=self.num_features ) self.discriminator = get_discriminator( - activation=activation, kernel_init=kernel_init, dropout_rate=dropout_rate, num_features=self.num_features, - num_additional_layers=num_additional_layers, cramer=cramer, features_to_tail=features_to_tail + activation=activation, + kernel_init=kernel_init, + dropout_rate=dropout_rate, + num_features=self.num_features, + num_additional_layers=num_additional_layers, + cramer=cramer, + features_to_tail=features_to_tail, ) self.step_counter = tf.Variable(0, dtype='int32', trainable=False) -# # compile the models with an arbitrary loss func for serializablility -# self.generator.compile(optimizer=self.gen_opt, -# loss='mean_squared_error') -# self.discriminator.compile(optimizer=self.disc_opt, -# loss='mean_squared_error') + # # compile the models with an arbitrary loss func for serializablility + # self.generator.compile(optimizer=self.gen_opt, + # 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) @@ -167,9 +178,7 @@ def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uni def make_fake(self, features): size = tf.shape(features)[0] latent_input = tf.random.normal(shape=(size, self.latent_dim), dtype='float32') - return self.generator( - tf.concat([_f(features), latent_input], axis=-1) - ) + return self.generator(tf.concat([_f(features), latent_input], axis=-1)) def gradient_penalty(self, features, real, fake): alpha = tf.random.uniform(shape=[len(real), 1, 1]) @@ -177,20 +186,20 @@ def gradient_penalty(self, features, real, fake): with tf.GradientTape() as t: t.watch(interpolates) d_int = self.discriminator([_f(features), interpolates]) -# if self.cramer: -# d_fake = self.discriminator([_f(features), interpolates]) -# d_int = tf.norm(d_int - d_fake, axis=-1) + # if self.cramer: + # d_fake = self.discriminator([_f(features), interpolates]) + # d_int = tf.norm(d_int - d_fake, axis=-1) 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) + return tf.reduce_mean(tf.maximum(tf.norm(grads, axis=-1) - 1, 0) ** 2) def gradient_penalty_on_data(self, features, real): with tf.GradientTape() as t: t.watch(real) d_real = self.discriminator([_f(features), real]) -# if self.cramer: -# d_real = tf.norm(d_real, axis=-1) + # if self.cramer: + # d_real = tf.norm(d_real, axis=-1) grads = tf.reshape(t.gradient(d_real, real), [len(real), -1]) - return tf.reduce_mean(tf.reduce_sum(grads**2, axis=-1)) + return tf.reduce_mean(tf.reduce_sum(grads ** 2, axis=-1)) @tf.function def calculate_losses(self, feature_batch, target_batch): @@ -207,19 +216,9 @@ def calculate_losses(self, feature_batch, target_batch): d_loss = disc_loss_cramer(d_real, d_fake, d_fake_2) if self.gp_lambda > 0: - d_loss = ( - d_loss + - self.gradient_penalty( - feature_batch, target_batch, fake - ) * self.gp_lambda - ) + d_loss = d_loss + self.gradient_penalty(feature_batch, target_batch, fake) * self.gp_lambda if self.gpdata_lambda > 0: - d_loss = ( - d_loss + - self.gradient_penalty_on_data( - feature_batch, target_batch - ) * self.gpdata_lambda - ) + d_loss = d_loss + self.gradient_penalty_on_data(feature_batch, target_batch) * self.gpdata_lambda if not self.cramer: g_loss = gen_loss(d_real, d_fake) else: @@ -252,10 +251,7 @@ def gen_step(self, feature_batch, target_batch): @tf.function def training_step(self, feature_batch, target_batch): if self.stochastic_stepping: - if tf.random.uniform( - shape=[], dtype='int32', - maxval=self.num_disc_updates + 1 - ) == self.num_disc_updates: + if tf.random.uniform(shape=[], dtype='int32', maxval=self.num_disc_updates + 1) == self.num_disc_updates: result = self.gen_step(feature_batch, target_batch) else: result = self.disc_step(feature_batch, target_batch) diff --git a/legacy_code/models/baseline_v2_10x10.py b/legacy_code/models/baseline_v2_10x10.py index 3f95e0d..adf23ca 100644 --- a/legacy_code/models/baseline_v2_10x10.py +++ b/legacy_code/models/baseline_v2_10x10.py @@ -2,74 +2,80 @@ def get_generator(activation, kernel_init, num_features, latent_dim): - generator = tf.keras.Sequential([ - tf.keras.layers.Dense(units=64, activation=activation, input_shape=(num_features + latent_dim,)), - - tf.keras.layers.Reshape((4, 4, 4)), - - 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=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=8, kernel_size=3, padding='valid', activation=activation, - kernel_initializer=kernel_init), # 10x10 - tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, - kernel_initializer=kernel_init), - - tf.keras.layers.Reshape((10, 10)), - ], name='generator') + generator = tf.keras.Sequential( + [ + tf.keras.layers.Dense(units=64, activation=activation, input_shape=(num_features + latent_dim,)), + tf.keras.layers.Reshape((4, 4, 4)), + 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=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=8, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init + ), # 10x10 + tf.keras.layers.Conv2D( + filters=1, + kernel_size=1, + padding='valid', + activation=tf.keras.activations.relu, + kernel_initializer=kernel_init, + ), + tf.keras.layers.Reshape((10, 10)), + ], + name='generator', + ) return generator -def get_discriminator(activation, kernel_init, dropout_rate, num_features, num_additional_layers, cramer=False, - features_to_tail=False): +def get_discriminator( + activation, kernel_init, dropout_rate, num_features, num_additional_layers, cramer=False, features_to_tail=False +): input_img = tf.keras.Input(shape=(10, 10)) features_input = tf.keras.Input(shape=(num_features,)) img = tf.reshape(input_img, (-1, 10, 10, 1)) if features_to_tail: - features_tiled = tf.tile( - tf.reshape(features_input, (-1, 1, 1, num_features)), - (1, 10, 10, 1) - ) - img = tf.concat( - [img, features_tiled], - axis=-1 - ) - - discriminator_tail = tf.keras.Sequential([ - 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(), # 4x4 - - 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.MaxPool2D(), # 2x2 - - 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.Reshape((64,)) - ], name='discriminator_tail') + features_tiled = tf.tile(tf.reshape(features_input, (-1, 1, 1, num_features)), (1, 10, 10, 1)) + img = tf.concat([img, features_tiled], axis=-1) + + discriminator_tail = tf.keras.Sequential( + [ + 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(), # 4x4 + 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.MaxPool2D(), # 2x2 + 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.Reshape((64,)), + ], + name='discriminator_tail', + ) head_input = tf.keras.layers.Concatenate()([features_input, discriminator_tail(img)]) @@ -84,19 +90,14 @@ def get_discriminator(activation, kernel_init, dropout_rate, num_features, num_a ] discriminator_head = tf.keras.Sequential( - head_layers + [tf.keras.layers.Dense(units=1 if not cramer else 256, - activation=None)], - name='discriminator_head' + head_layers + [tf.keras.layers.Dense(units=1 if not cramer else 256, activation=None)], + name='discriminator_head', ) inputs = [features_input, input_img] outputs = discriminator_head(head_input) - discriminator = tf.keras.Model( - inputs=inputs, - outputs=outputs, - name='discriminator' - ) + discriminator = tf.keras.Model(inputs=inputs, outputs=outputs, name='discriminator') return discriminator @@ -111,10 +112,10 @@ def gen_loss(d_real, d_fake): def disc_loss_cramer(d_real, d_fake, d_fake_2): return -tf.reduce_mean( - tf.norm(d_real - d_fake, axis=-1) + - tf.norm(d_fake_2, axis=-1) - - tf.norm(d_fake - d_fake_2, axis=-1) - - tf.norm(d_real, axis=-1) + tf.norm(d_real - d_fake, axis=-1) + + tf.norm(d_fake_2, axis=-1) + - tf.norm(d_fake - d_fake_2, axis=-1) + - tf.norm(d_real, axis=-1) ) @@ -123,10 +124,22 @@ def gen_loss_cramer(d_real, d_fake, d_fake_2): class BaselineModel10x10: - def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uniform', - dropout_rate=0.2, lr=1e-4, latent_dim=32, gp_lambda=10., num_disc_updates=3, - num_features=1, gpdata_lambda=0., num_additional_layers=0, cramer=False, - features_to_tail=False, stochastic_stepping=False): + def __init__( + self, + activation=tf.keras.activations.relu, + kernel_init='glorot_uniform', + dropout_rate=0.2, + lr=1e-4, + latent_dim=32, + gp_lambda=10.0, + num_disc_updates=3, + num_features=1, + gpdata_lambda=0.0, + num_additional_layers=0, + cramer=False, + features_to_tail=False, + stochastic_stepping=False, + ): self.disc_opt = tf.keras.optimizers.RMSprop(lr) self.gen_opt = tf.keras.optimizers.RMSprop(lr) self.latent_dim = latent_dim @@ -141,8 +154,13 @@ def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uni activation=activation, kernel_init=kernel_init, latent_dim=latent_dim, num_features=num_features ) self.discriminator = get_discriminator( - activation=activation, kernel_init=kernel_init, dropout_rate=dropout_rate, num_features=num_features, - num_additional_layers=num_additional_layers, cramer=cramer, features_to_tail=features_to_tail + activation=activation, + kernel_init=kernel_init, + dropout_rate=dropout_rate, + num_features=num_features, + num_additional_layers=num_additional_layers, + cramer=cramer, + features_to_tail=features_to_tail, ) self.step_counter = tf.Variable(0, dtype='int32', trainable=False) @@ -150,9 +168,7 @@ def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uni def make_fake(self, features): size = tf.shape(features)[0] latent_input = tf.random.normal(shape=(size, self.latent_dim), dtype='float32') - return self.generator( - tf.concat([features, latent_input], axis=-1) - ) + return self.generator(tf.concat([features, latent_input], axis=-1)) def gradient_penalty(self, features, real, fake): alpha = tf.random.uniform(shape=[len(real), 1, 1]) @@ -160,20 +176,20 @@ def gradient_penalty(self, features, real, fake): with tf.GradientTape() as t: t.watch(interpolates) d_int = self.discriminator([features, interpolates]) -# if self.cramer: -# d_fake = self.discriminator([features, interpolates]) -# d_int = tf.norm(d_int - d_fake, axis=-1) + # if self.cramer: + # d_fake = self.discriminator([features, interpolates]) + # d_int = tf.norm(d_int - d_fake, axis=-1) 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) + return tf.reduce_mean(tf.maximum(tf.norm(grads, axis=-1) - 1, 0) ** 2) def gradient_penalty_on_data(self, features, real): with tf.GradientTape() as t: t.watch(real) d_real = self.discriminator([features, real]) -# if self.cramer: -# d_real = tf.norm(d_real, axis=-1) + # if self.cramer: + # d_real = tf.norm(d_real, axis=-1) grads = tf.reshape(t.gradient(d_real, real), [len(real), -1]) - return tf.reduce_mean(tf.reduce_sum(grads**2, axis=-1)) + return tf.reduce_mean(tf.reduce_sum(grads ** 2, axis=-1)) @tf.function def calculate_losses(self, feature_batch, target_batch): @@ -190,19 +206,9 @@ def calculate_losses(self, feature_batch, target_batch): d_loss = disc_loss_cramer(d_real, d_fake, d_fake_2) if self.gp_lambda > 0: - d_loss = ( - d_loss + - self.gradient_penalty( - feature_batch, target_batch, fake - ) * self.gp_lambda - ) + d_loss = d_loss + self.gradient_penalty(feature_batch, target_batch, fake) * self.gp_lambda if self.gpdata_lambda > 0: - d_loss = ( - d_loss + - self.gradient_penalty_on_data( - feature_batch, target_batch - ) * self.gpdata_lambda - ) + d_loss = d_loss + self.gradient_penalty_on_data(feature_batch, target_batch) * self.gpdata_lambda if not self.cramer: g_loss = gen_loss(d_real, d_fake) else: @@ -235,10 +241,7 @@ def gen_step(self, feature_batch, target_batch): @tf.function def training_step(self, feature_batch, target_batch): if self.stochastic_stepping: - if tf.random.uniform( - shape=[], dtype='int32', - maxval=self.num_disc_updates + 1 - ) == self.num_disc_updates: + if tf.random.uniform(shape=[], dtype='int32', maxval=self.num_disc_updates + 1) == self.num_disc_updates: result = self.gen_step(feature_batch, target_batch) else: result = self.disc_step(feature_batch, target_batch) diff --git a/legacy_code/models/baseline_v3_6x15.py b/legacy_code/models/baseline_v3_6x15.py index 6548dfc..0486529 100644 --- a/legacy_code/models/baseline_v3_6x15.py +++ b/legacy_code/models/baseline_v3_6x15.py @@ -8,9 +8,7 @@ def preprocess_features(features): # dip_angle [-60, 60] # drift_length [35, 290] bin_fraction = features[:, -1:] % 1 - features = ( - features - tf.constant([[0., 0., 162.5]]) - ) / tf.constant([[20., 60., 127.5]]) + features = (features - tf.constant([[0.0, 0.0, 162.5]])) / tf.constant([[20.0, 60.0, 127.5]]) return tf.concat([features, bin_fraction], axis=-1) @@ -18,76 +16,83 @@ def preprocess_features(features): def get_generator(activation, kernel_init, num_features, latent_dim): - generator = tf.keras.Sequential([ - tf.keras.layers.Dense(units=64, activation=activation, input_shape=(num_features + latent_dim,)), - - tf.keras.layers.Reshape((4, 4, 4)), - - 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(size=(1, 2)), # 4x8 - - 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='same', activation=activation, - kernel_initializer=kernel_init), - tf.keras.layers.UpSampling2D(), # 8x16 - - tf.keras.layers.Conv2D(filters=8, kernel_size=(1, 3), padding='same', activation=activation, - kernel_initializer=kernel_init), - tf.keras.layers.Conv2D(filters=8, kernel_size=(3, 2), padding='valid', activation=activation, - kernel_initializer=kernel_init), # 6x15 - tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, - kernel_initializer=kernel_init), - - tf.keras.layers.Reshape((6, 15)), - ], name='generator') + generator = tf.keras.Sequential( + [ + tf.keras.layers.Dense(units=64, activation=activation, input_shape=(num_features + latent_dim,)), + tf.keras.layers.Reshape((4, 4, 4)), + 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(size=(1, 2)), # 4x8 + 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='same', activation=activation, kernel_initializer=kernel_init + ), + tf.keras.layers.UpSampling2D(), # 8x16 + tf.keras.layers.Conv2D( + filters=8, kernel_size=(1, 3), padding='same', activation=activation, kernel_initializer=kernel_init + ), + tf.keras.layers.Conv2D( + filters=8, kernel_size=(3, 2), padding='valid', activation=activation, kernel_initializer=kernel_init + ), # 6x15 + tf.keras.layers.Conv2D( + filters=1, + kernel_size=1, + padding='valid', + activation=tf.keras.activations.relu, + kernel_initializer=kernel_init, + ), + tf.keras.layers.Reshape((6, 15)), + ], + name='generator', + ) return generator -def get_discriminator(activation, kernel_init, dropout_rate, num_features, num_additional_layers, cramer=False, - features_to_tail=False): +def get_discriminator( + activation, kernel_init, dropout_rate, num_features, num_additional_layers, cramer=False, features_to_tail=False +): input_img = tf.keras.Input(shape=(6, 15)) features_input = tf.keras.Input(shape=(num_features,)) img = tf.reshape(input_img, (-1, 6, 15, 1)) if features_to_tail: - features_tiled = tf.tile( - tf.reshape(features_input, (-1, 1, 1, num_features)), - (1, 6, 15, 1) - ) - img = tf.concat( - [img, features_tiled], - axis=-1 - ) - - discriminator_tail = tf.keras.Sequential([ - 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=(1, 4), padding='valid', activation=activation, - kernel_initializer=kernel_init), # 6x12 - tf.keras.layers.Dropout(dropout_rate), - - tf.keras.layers.MaxPool2D(pool_size=(1, 2)), # 6x6 - - 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.MaxPool2D(), # 3x3 - - tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation, - kernel_initializer=kernel_init), # 1x1 - tf.keras.layers.Dropout(dropout_rate), - - tf.keras.layers.Reshape((64,)) - ], name='discriminator_tail') + features_tiled = tf.tile(tf.reshape(features_input, (-1, 1, 1, num_features)), (1, 6, 15, 1)) + img = tf.concat([img, features_tiled], axis=-1) + + discriminator_tail = tf.keras.Sequential( + [ + 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=(1, 4), padding='valid', activation=activation, kernel_initializer=kernel_init + ), # 6x12 + tf.keras.layers.Dropout(dropout_rate), + tf.keras.layers.MaxPool2D(pool_size=(1, 2)), # 6x6 + 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.MaxPool2D(), # 3x3 + tf.keras.layers.Conv2D( + filters=64, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init + ), # 1x1 + tf.keras.layers.Dropout(dropout_rate), + tf.keras.layers.Reshape((64,)), + ], + name='discriminator_tail', + ) head_input = tf.keras.layers.Concatenate()([features_input, discriminator_tail(img)]) @@ -102,19 +107,14 @@ def get_discriminator(activation, kernel_init, dropout_rate, num_features, num_a ] discriminator_head = tf.keras.Sequential( - head_layers + [tf.keras.layers.Dense(units=1 if not cramer else 256, - activation=None)], - name='discriminator_head' + head_layers + [tf.keras.layers.Dense(units=1 if not cramer else 256, activation=None)], + name='discriminator_head', ) inputs = [features_input, input_img] outputs = discriminator_head(head_input) - discriminator = tf.keras.Model( - inputs=inputs, - outputs=outputs, - name='discriminator' - ) + discriminator = tf.keras.Model(inputs=inputs, outputs=outputs, name='discriminator') return discriminator @@ -129,10 +129,10 @@ def gen_loss(d_real, d_fake): def disc_loss_cramer(d_real, d_fake, d_fake_2): return -tf.reduce_mean( - tf.norm(d_real - d_fake, axis=-1) + - tf.norm(d_fake_2, axis=-1) - - tf.norm(d_fake - d_fake_2, axis=-1) - - tf.norm(d_real, axis=-1) + tf.norm(d_real - d_fake, axis=-1) + + tf.norm(d_fake_2, axis=-1) + - tf.norm(d_fake - d_fake_2, axis=-1) + - tf.norm(d_real, axis=-1) ) @@ -141,10 +141,21 @@ def gen_loss_cramer(d_real, d_fake, d_fake_2): class BaselineModel_6x15: - def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uniform', - dropout_rate=0.02, lr=1e-4, latent_dim=32, gp_lambda=10., num_disc_updates=8, - gpdata_lambda=0., num_additional_layers=0, cramer=False, - features_to_tail=True, stochastic_stepping=True): + def __init__( + self, + activation=tf.keras.activations.relu, + kernel_init='glorot_uniform', + dropout_rate=0.02, + lr=1e-4, + latent_dim=32, + gp_lambda=10.0, + num_disc_updates=8, + gpdata_lambda=0.0, + num_additional_layers=0, + cramer=False, + features_to_tail=True, + stochastic_stepping=True, + ): self.disc_opt = tf.keras.optimizers.RMSprop(lr) self.gen_opt = tf.keras.optimizers.RMSprop(lr) self.latent_dim = latent_dim @@ -159,25 +170,28 @@ def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uni activation=activation, kernel_init=kernel_init, latent_dim=latent_dim, num_features=self.num_features ) self.discriminator = get_discriminator( - activation=activation, kernel_init=kernel_init, dropout_rate=dropout_rate, num_features=self.num_features, - num_additional_layers=num_additional_layers, cramer=cramer, features_to_tail=features_to_tail + activation=activation, + kernel_init=kernel_init, + dropout_rate=dropout_rate, + num_features=self.num_features, + num_additional_layers=num_additional_layers, + cramer=cramer, + features_to_tail=features_to_tail, ) self.step_counter = tf.Variable(0, dtype='int32', trainable=False) -# # compile the models with an arbitrary loss func for serializablility -# self.generator.compile(optimizer=self.gen_opt, -# loss='mean_squared_error') -# self.discriminator.compile(optimizer=self.disc_opt, -# loss='mean_squared_error') + # # compile the models with an arbitrary loss func for serializablility + # self.generator.compile(optimizer=self.gen_opt, + # loss='mean_squared_error') + # self.discriminator.compile(optimizer=self.disc_opt, + # loss='mean_squared_error') @tf.function def make_fake(self, features): size = tf.shape(features)[0] latent_input = tf.random.normal(shape=(size, self.latent_dim), dtype='float32') - return self.generator( - tf.concat([_f(features), latent_input], axis=-1) - ) + return self.generator(tf.concat([_f(features), latent_input], axis=-1)) def gradient_penalty(self, features, real, fake): alpha = tf.random.uniform(shape=[len(real), 1, 1]) @@ -185,20 +199,20 @@ def gradient_penalty(self, features, real, fake): with tf.GradientTape() as t: t.watch(interpolates) d_int = self.discriminator([_f(features), interpolates]) -# if self.cramer: -# d_fake = self.discriminator([_f(features), interpolates]) -# d_int = tf.norm(d_int - d_fake, axis=-1) + # if self.cramer: + # d_fake = self.discriminator([_f(features), interpolates]) + # d_int = tf.norm(d_int - d_fake, axis=-1) 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) + return tf.reduce_mean(tf.maximum(tf.norm(grads, axis=-1) - 1, 0) ** 2) def gradient_penalty_on_data(self, features, real): with tf.GradientTape() as t: t.watch(real) d_real = self.discriminator([_f(features), real]) -# if self.cramer: -# d_real = tf.norm(d_real, axis=-1) + # if self.cramer: + # d_real = tf.norm(d_real, axis=-1) grads = tf.reshape(t.gradient(d_real, real), [len(real), -1]) - return tf.reduce_mean(tf.reduce_sum(grads**2, axis=-1)) + return tf.reduce_mean(tf.reduce_sum(grads ** 2, axis=-1)) @tf.function def calculate_losses(self, feature_batch, target_batch): @@ -215,19 +229,9 @@ def calculate_losses(self, feature_batch, target_batch): d_loss = disc_loss_cramer(d_real, d_fake, d_fake_2) if self.gp_lambda > 0: - d_loss = ( - d_loss + - self.gradient_penalty( - feature_batch, target_batch, fake - ) * self.gp_lambda - ) + d_loss = d_loss + self.gradient_penalty(feature_batch, target_batch, fake) * self.gp_lambda if self.gpdata_lambda > 0: - d_loss = ( - d_loss + - self.gradient_penalty_on_data( - feature_batch, target_batch - ) * self.gpdata_lambda - ) + d_loss = d_loss + self.gradient_penalty_on_data(feature_batch, target_batch) * self.gpdata_lambda if not self.cramer: g_loss = gen_loss(d_real, d_fake) else: @@ -260,10 +264,7 @@ def gen_step(self, feature_batch, target_batch): @tf.function def training_step(self, feature_batch, target_batch): if self.stochastic_stepping: - if tf.random.uniform( - shape=[], dtype='int32', - maxval=self.num_disc_updates + 1 - ) == self.num_disc_updates: + if tf.random.uniform(shape=[], dtype='int32', maxval=self.num_disc_updates + 1) == self.num_disc_updates: result = self.gen_step(feature_batch, target_batch) else: result = self.disc_step(feature_batch, target_batch) diff --git a/legacy_code/models/baseline_v4_8x16.py b/legacy_code/models/baseline_v4_8x16.py index c992ea2..2e37e94 100644 --- a/legacy_code/models/baseline_v4_8x16.py +++ b/legacy_code/models/baseline_v4_8x16.py @@ -9,9 +9,7 @@ def preprocess_features(features): # drift_length [35, 290] # pad_coordinate [40-something, 40-something] bin_fractions = features[:, -2:] % 1 - features = ( - features[:, :3] - tf.constant([[0., 0., 162.5]]) - ) / tf.constant([[20., 60., 127.5]]) + features = (features[:, :3] - tf.constant([[0.0, 0.0, 162.5]])) / tf.constant([[20.0, 60.0, 127.5]]) return tf.concat([features, bin_fractions], axis=-1) @@ -19,83 +17,91 @@ def preprocess_features(features): def get_generator(activation, kernel_init, num_features, latent_dim): - generator = tf.keras.Sequential([ - tf.keras.layers.Dense(units=64, activation=activation, input_shape=(num_features + latent_dim,)), - - tf.keras.layers.Reshape((2, 4, 8)), - - tf.keras.layers.Conv2D(filters=32, kernel_size=2, padding='same', activation=activation, - kernel_initializer=kernel_init), - tf.keras.layers.Conv2D(filters=32, kernel_size=2, padding='same', activation=activation, - kernel_initializer=kernel_init), # 2x4 - tf.keras.layers.UpSampling2D(), # 4x8 - - tf.keras.layers.Conv2D(filters=16, kernel_size=(2, 3), padding='valid', activation=activation, - kernel_initializer=kernel_init), # 3x6 - tf.keras.layers.Conv2D(filters=16, kernel_size=(1, 2), padding='valid', activation=activation, - kernel_initializer=kernel_init), # 3x5 - tf.keras.layers.UpSampling2D(), # 6x10 - - tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='same', activation=activation, - kernel_initializer=kernel_init), - tf.keras.layers.Conv2D(filters=8, kernel_size=2, padding='valid', activation=activation, - kernel_initializer=kernel_init), # 5x9 - tf.keras.layers.UpSampling2D(), # 10x18 - - tf.keras.layers.Conv2D(filters=8, kernel_size=3, padding='valid', activation=activation, - kernel_initializer=kernel_init), # 8x16 - tf.keras.layers.Conv2D(filters=1, kernel_size=1, padding='valid', activation=tf.keras.activations.relu, - kernel_initializer=kernel_init), - - tf.keras.layers.Reshape((8, 16)), - ], name='generator') + generator = tf.keras.Sequential( + [ + tf.keras.layers.Dense(units=64, activation=activation, input_shape=(num_features + latent_dim,)), + tf.keras.layers.Reshape((2, 4, 8)), + tf.keras.layers.Conv2D( + filters=32, kernel_size=2, padding='same', activation=activation, kernel_initializer=kernel_init + ), + tf.keras.layers.Conv2D( + filters=32, kernel_size=2, padding='same', activation=activation, kernel_initializer=kernel_init + ), # 2x4 + tf.keras.layers.UpSampling2D(), # 4x8 + tf.keras.layers.Conv2D( + filters=16, kernel_size=(2, 3), padding='valid', activation=activation, kernel_initializer=kernel_init + ), # 3x6 + tf.keras.layers.Conv2D( + filters=16, kernel_size=(1, 2), padding='valid', activation=activation, kernel_initializer=kernel_init + ), # 3x5 + tf.keras.layers.UpSampling2D(), # 6x10 + tf.keras.layers.Conv2D( + filters=8, kernel_size=3, padding='same', activation=activation, kernel_initializer=kernel_init + ), + tf.keras.layers.Conv2D( + filters=8, kernel_size=2, padding='valid', activation=activation, kernel_initializer=kernel_init + ), # 5x9 + tf.keras.layers.UpSampling2D(), # 10x18 + tf.keras.layers.Conv2D( + filters=8, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init + ), # 8x16 + tf.keras.layers.Conv2D( + filters=1, + kernel_size=1, + padding='valid', + activation=tf.keras.activations.relu, + kernel_initializer=kernel_init, + ), + tf.keras.layers.Reshape((8, 16)), + ], + name='generator', + ) return generator -def get_discriminator(activation, kernel_init, dropout_rate, num_features, num_additional_layers, cramer=False, - features_to_tail=False): +def get_discriminator( + activation, kernel_init, dropout_rate, num_features, num_additional_layers, cramer=False, features_to_tail=False +): input_img = tf.keras.Input(shape=(8, 16)) features_input = tf.keras.Input(shape=(num_features,)) img = tf.reshape(input_img, (-1, 8, 16, 1)) if features_to_tail: - features_tiled = tf.tile( - tf.reshape(features_input, (-1, 1, 1, num_features)), - (1, 8, 16, 1) - ) - img = tf.concat( - [img, features_tiled], - axis=-1 - ) - - discriminator_tail = tf.keras.Sequential([ - 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='same', activation=activation, - kernel_initializer=kernel_init), - tf.keras.layers.Dropout(dropout_rate), - - tf.keras.layers.MaxPool2D(pool_size=(1, 2)), # 8x8 - - 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.MaxPool2D(), # 4x4 - - tf.keras.layers.Conv2D(filters=64, kernel_size=3, padding='valid', activation=activation, - kernel_initializer=kernel_init), # 2x2 - 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.Reshape((64,)) - ], name='discriminator_tail') + features_tiled = tf.tile(tf.reshape(features_input, (-1, 1, 1, num_features)), (1, 8, 16, 1)) + img = tf.concat([img, features_tiled], axis=-1) + + discriminator_tail = tf.keras.Sequential( + [ + 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='same', activation=activation, kernel_initializer=kernel_init + ), + tf.keras.layers.Dropout(dropout_rate), + tf.keras.layers.MaxPool2D(pool_size=(1, 2)), # 8x8 + 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.MaxPool2D(), # 4x4 + tf.keras.layers.Conv2D( + filters=64, kernel_size=3, padding='valid', activation=activation, kernel_initializer=kernel_init + ), # 2x2 + 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.Reshape((64,)), + ], + name='discriminator_tail', + ) head_input = tf.keras.layers.Concatenate()([features_input, discriminator_tail(img)]) @@ -110,19 +116,14 @@ def get_discriminator(activation, kernel_init, dropout_rate, num_features, num_a ] discriminator_head = tf.keras.Sequential( - head_layers + [tf.keras.layers.Dense(units=1 if not cramer else 256, - activation=None)], - name='discriminator_head' + head_layers + [tf.keras.layers.Dense(units=1 if not cramer else 256, activation=None)], + name='discriminator_head', ) inputs = [features_input, input_img] outputs = discriminator_head(head_input) - discriminator = tf.keras.Model( - inputs=inputs, - outputs=outputs, - name='discriminator' - ) + discriminator = tf.keras.Model(inputs=inputs, outputs=outputs, name='discriminator') return discriminator @@ -137,10 +138,10 @@ def gen_loss(d_real, d_fake): def disc_loss_cramer(d_real, d_fake, d_fake_2): return -tf.reduce_mean( - tf.norm(d_real - d_fake, axis=-1) + - tf.norm(d_fake_2, axis=-1) - - tf.norm(d_fake - d_fake_2, axis=-1) - - tf.norm(d_real, axis=-1) + tf.norm(d_real - d_fake, axis=-1) + + tf.norm(d_fake_2, axis=-1) + - tf.norm(d_fake - d_fake_2, axis=-1) + - tf.norm(d_real, axis=-1) ) @@ -149,10 +150,21 @@ def gen_loss_cramer(d_real, d_fake, d_fake_2): class BaselineModel_8x16: - def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uniform', - dropout_rate=0.02, lr=1e-4, latent_dim=32, gp_lambda=10., num_disc_updates=8, - gpdata_lambda=0., num_additional_layers=0, cramer=False, - features_to_tail=True, stochastic_stepping=True): + def __init__( + self, + activation=tf.keras.activations.relu, + kernel_init='glorot_uniform', + dropout_rate=0.02, + lr=1e-4, + latent_dim=32, + gp_lambda=10.0, + num_disc_updates=8, + gpdata_lambda=0.0, + num_additional_layers=0, + cramer=False, + features_to_tail=True, + stochastic_stepping=True, + ): self.disc_opt = tf.keras.optimizers.RMSprop(lr) self.gen_opt = tf.keras.optimizers.RMSprop(lr) self.latent_dim = latent_dim @@ -167,25 +179,28 @@ def __init__(self, activation=tf.keras.activations.relu, kernel_init='glorot_uni activation=activation, kernel_init=kernel_init, latent_dim=latent_dim, num_features=self.num_features ) self.discriminator = get_discriminator( - activation=activation, kernel_init=kernel_init, dropout_rate=dropout_rate, num_features=self.num_features, - num_additional_layers=num_additional_layers, cramer=cramer, features_to_tail=features_to_tail + activation=activation, + kernel_init=kernel_init, + dropout_rate=dropout_rate, + num_features=self.num_features, + num_additional_layers=num_additional_layers, + cramer=cramer, + features_to_tail=features_to_tail, ) self.step_counter = tf.Variable(0, dtype='int32', trainable=False) -# # compile the models with an arbitrary loss func for serializablility -# self.generator.compile(optimizer=self.gen_opt, -# loss='mean_squared_error') -# self.discriminator.compile(optimizer=self.disc_opt, -# loss='mean_squared_error') + # # compile the models with an arbitrary loss func for serializablility + # self.generator.compile(optimizer=self.gen_opt, + # loss='mean_squared_error') + # self.discriminator.compile(optimizer=self.disc_opt, + # loss='mean_squared_error') @tf.function def make_fake(self, features): size = tf.shape(features)[0] latent_input = tf.random.normal(shape=(size, self.latent_dim), dtype='float32') - return self.generator( - tf.concat([_f(features), latent_input], axis=-1) - ) + return self.generator(tf.concat([_f(features), latent_input], axis=-1)) def gradient_penalty(self, features, real, fake): alpha = tf.random.uniform(shape=[len(real), 1, 1]) @@ -193,20 +208,20 @@ def gradient_penalty(self, features, real, fake): with tf.GradientTape() as t: t.watch(interpolates) d_int = self.discriminator([_f(features), interpolates]) -# if self.cramer: -# d_fake = self.discriminator([_f(features), interpolates]) -# d_int = tf.norm(d_int - d_fake, axis=-1) + # if self.cramer: + # d_fake = self.discriminator([_f(features), interpolates]) + # d_int = tf.norm(d_int - d_fake, axis=-1) 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) + return tf.reduce_mean(tf.maximum(tf.norm(grads, axis=-1) - 1, 0) ** 2) def gradient_penalty_on_data(self, features, real): with tf.GradientTape() as t: t.watch(real) d_real = self.discriminator([_f(features), real]) -# if self.cramer: -# d_real = tf.norm(d_real, axis=-1) + # if self.cramer: + # d_real = tf.norm(d_real, axis=-1) grads = tf.reshape(t.gradient(d_real, real), [len(real), -1]) - return tf.reduce_mean(tf.reduce_sum(grads**2, axis=-1)) + return tf.reduce_mean(tf.reduce_sum(grads ** 2, axis=-1)) @tf.function def calculate_losses(self, feature_batch, target_batch): @@ -223,19 +238,9 @@ def calculate_losses(self, feature_batch, target_batch): d_loss = disc_loss_cramer(d_real, d_fake, d_fake_2) if self.gp_lambda > 0: - d_loss = ( - d_loss + - self.gradient_penalty( - feature_batch, target_batch, fake - ) * self.gp_lambda - ) + d_loss = d_loss + self.gradient_penalty(feature_batch, target_batch, fake) * self.gp_lambda if self.gpdata_lambda > 0: - d_loss = ( - d_loss + - self.gradient_penalty_on_data( - feature_batch, target_batch - ) * self.gpdata_lambda - ) + d_loss = d_loss + self.gradient_penalty_on_data(feature_batch, target_batch) * self.gpdata_lambda if not self.cramer: g_loss = gen_loss(d_real, d_fake) else: @@ -268,10 +273,7 @@ def gen_step(self, feature_batch, target_batch): @tf.function def training_step(self, feature_batch, target_batch): if self.stochastic_stepping: - if tf.random.uniform( - shape=[], dtype='int32', - maxval=self.num_disc_updates + 1 - ) == self.num_disc_updates: + if tf.random.uniform(shape=[], dtype='int32', maxval=self.num_disc_updates + 1) == self.num_disc_updates: result = self.gen_step(feature_batch, target_batch) else: result = self.disc_step(feature_batch, target_batch) diff --git a/legacy_code/test_script_data_v0.py b/legacy_code/test_script_data_v0.py index c0cbf0c..77dbd72 100644 --- a/legacy_code/test_script_data_v0.py +++ b/legacy_code/test_script_data_v0.py @@ -36,7 +36,7 @@ def write_hist_summary(step): real = unscale(X_test) gen = unscale(gen_scaled) gen[gen < 0] = 0 - gen30 = np.where(gen < 30., 0, gen) + gen30 = np.where(gen < 30.0, 0, gen) images = make_metric_plots(real, gen) images30 = make_metric_plots(real, gen30) @@ -56,18 +56,26 @@ def write_hist_summary(step): def save_model(step): if step % 50 == 0: - baseline_10x15.generator .save(str(model_path.joinpath("generator_{:05d}.h5" .format(step)))) + baseline_10x15.generator.save(str(model_path.joinpath("generator_{:05d}.h5".format(step)))) baseline_10x15.discriminator.save(str(model_path.joinpath("discriminator_{:05d}.h5".format(step)))) def schedule_lr(step): baseline_10x15.disc_opt.lr.assign(baseline_10x15.disc_opt.lr * 0.998) - baseline_10x15.gen_opt .lr.assign(baseline_10x15.gen_opt .lr * 0.998) + baseline_10x15.gen_opt.lr.assign(baseline_10x15.gen_opt.lr * 0.998) with writer_val.as_default(): tf.summary.scalar("discriminator learning rate", baseline_10x15.disc_opt.lr, step) tf.summary.scalar("generator learning rate", baseline_10x15.gen_opt.lr, step) -training.train(X_train, X_test, baseline_10x15.training_step, baseline_10x15.calculate_losses, 10000, 32, - train_writer=writer_train, val_writer=writer_val, - callbacks=[write_hist_summary, save_model, schedule_lr]) +training.train( + X_train, + X_test, + baseline_10x15.training_step, + baseline_10x15.calculate_losses, + 10000, + 32, + train_writer=writer_train, + val_writer=writer_val, + callbacks=[write_hist_summary, save_model, schedule_lr], +) diff --git a/legacy_code/test_script_data_v1.py b/legacy_code/test_script_data_v1.py index 4f9e9e7..e2a1b55 100644 --- a/legacy_code/test_script_data_v1.py +++ b/legacy_code/test_script_data_v1.py @@ -38,8 +38,9 @@ def main(): model_path = Path('saved_models') / args.checkpoint_name model_path.mkdir(parents=True) - model = BaselineModel10x10(kernel_init=args.kernel_init, lr=args.lr, - num_disc_updates=args.num_disc_updates, latent_dim=args.latent_dim) + model = BaselineModel10x10( + kernel_init=args.kernel_init, lr=args.lr, num_disc_updates=args.num_disc_updates, latent_dim=args.latent_dim + ) def save_model(step): if step % args.save_every == 0: @@ -65,7 +66,7 @@ def write_hist_summary(step): real = unscale(X_test) gen = unscale(gen_scaled) gen[gen < 0] = 0 - gen1 = np.where(gen < 1., 0, gen) + gen1 = np.where(gen < 1.0, 0, gen) images = make_metric_plots(real, gen) images1 = make_metric_plots(real, gen1) @@ -85,9 +86,17 @@ def schedule_lr(step): tf.summary.scalar("discriminator learning rate", model.disc_opt.lr, step) tf.summary.scalar("generator learning rate", model.gen_opt.lr, step) - train(X_train, X_test, model.training_step, model.calculate_losses, args.num_epochs, args.batch_size, - train_writer=writer_train, val_writer=writer_val, - callbacks=[write_hist_summary, save_model, schedule_lr]) + train( + X_train, + X_test, + model.training_step, + model.calculate_losses, + args.num_epochs, + args.batch_size, + train_writer=writer_train, + val_writer=writer_val, + callbacks=[write_hist_summary, save_model, schedule_lr], + ) if __name__ == '__main__': diff --git a/legacy_code/test_script_data_v1_normed.py b/legacy_code/test_script_data_v1_normed.py index ca14403..a83d391 100644 --- a/legacy_code/test_script_data_v1_normed.py +++ b/legacy_code/test_script_data_v1_normed.py @@ -59,15 +59,23 @@ def save_model(step): def schedule_lr(step): baseline_10x10_normed.disc_opt.lr.assign(baseline_10x10_normed.disc_opt.lr * 0.998) - baseline_10x10_normed.gen_opt .lr.assign(baseline_10x10_normed.gen_opt .lr * 0.998) + baseline_10x10_normed.gen_opt.lr.assign(baseline_10x10_normed.gen_opt.lr * 0.998) with writer_val.as_default(): tf.summary.scalar("discriminator learning rate", baseline_10x10_normed.disc_opt.lr, step) - tf.summary.scalar("generator learning rate", baseline_10x10_normed.gen_opt .lr, step) + tf.summary.scalar("generator learning rate", baseline_10x10_normed.gen_opt.lr, step) baseline_10x10_normed.disc_opt.lr.assign(0.0001) -baseline_10x10_normed.gen_opt .lr.assign(0.0001) - -training.train(X_train, X_test, baseline_10x10_normed.training_step, baseline_10x10_normed.calculate_losses, 10000, 32, - train_writer=writer_train, val_writer=writer_val, - callbacks=[write_hist_summary, save_model, schedule_lr]) +baseline_10x10_normed.gen_opt.lr.assign(0.0001) + +training.train( + X_train, + X_test, + baseline_10x10_normed.training_step, + baseline_10x10_normed.calculate_losses, + 10000, + 32, + train_writer=writer_train, + val_writer=writer_val, + callbacks=[write_hist_summary, save_model, schedule_lr], +) diff --git a/legacy_code/test_script_data_v2.py b/legacy_code/test_script_data_v2.py index 836d184..d500df5 100644 --- a/legacy_code/test_script_data_v2.py +++ b/legacy_code/test_script_data_v2.py @@ -27,8 +27,8 @@ def main(): parser.add_argument('--latent_dim', type=int, default=32, required=False) parser.add_argument('--gpu_num', type=str, required=False) parser.add_argument('--kernel_init', type=str, default='glorot_uniform', required=False) - parser.add_argument('--gp_lambda', type=float, default=10., required=False) - parser.add_argument('--gpdata_lambda', type=float, default=0., required=False) + parser.add_argument('--gp_lambda', type=float, default=10.0, required=False) + parser.add_argument('--gpdata_lambda', type=float, default=0.0, required=False) parser.add_argument('--angles_to_radians', action='store_true') parser.add_argument('--num_additional_disc_layers', type=int, default=0, required=False) parser.add_argument('--cramer_gan', action='store_true', default=False) @@ -41,9 +41,8 @@ def main(): args = parser.parse_args() - assert ( - (args.feature_noise_power is None) == - (args.feature_noise_decay is None) + assert (args.feature_noise_power is None) == ( + args.feature_noise_decay is None ), 'Noise power and decay must be both provided' print("") @@ -82,32 +81,33 @@ def main(): f.write('\n') f.write(f_in.read()) - model = BaselineModel10x10(kernel_init=args.kernel_init, lr=args.lr, - num_disc_updates=args.num_disc_updates, latent_dim=args.latent_dim, - gp_lambda=args.gp_lambda, gpdata_lambda=args.gpdata_lambda, - num_additional_layers=args.num_additional_disc_layers, - cramer=args.cramer_gan, features_to_tail=args.features_to_tail, - dropout_rate=args.dropout_rate, - stochastic_stepping=args.stochastic_stepping) + model = BaselineModel10x10( + kernel_init=args.kernel_init, + lr=args.lr, + num_disc_updates=args.num_disc_updates, + latent_dim=args.latent_dim, + gp_lambda=args.gp_lambda, + gpdata_lambda=args.gpdata_lambda, + num_additional_layers=args.num_additional_disc_layers, + cramer=args.cramer_gan, + features_to_tail=args.features_to_tail, + dropout_rate=args.dropout_rate, + stochastic_stepping=args.stochastic_stepping, + ) if args.prediction_only: + def epoch_from_name(name): - epoch, = re.findall(r'\d+', name) + (epoch,) = re.findall(r'\d+', name) return int(epoch) gen_checkpoints = model_path.glob("generator_*.h5") disc_checkpoints = model_path.glob("discriminator_*.h5") - latest_gen_checkpoint = max( - gen_checkpoints, - key=lambda path: epoch_from_name(path.stem) - ) - latest_disc_checkpoint = max( - disc_checkpoints, - key=lambda path: epoch_from_name(path.stem) - ) + latest_gen_checkpoint = max(gen_checkpoints, key=lambda path: epoch_from_name(path.stem)) + latest_disc_checkpoint = max(disc_checkpoints, key=lambda path: epoch_from_name(path.stem)) - assert ( - epoch_from_name(latest_gen_checkpoint.stem) == epoch_from_name(latest_disc_checkpoint.stem) + assert epoch_from_name(latest_gen_checkpoint.stem) == epoch_from_name( + latest_disc_checkpoint.stem ), "Latest disc and gen epochs differ" print(f'Loading generator weights from {str(latest_gen_checkpoint)}') @@ -145,18 +145,18 @@ def get_images(return_raw_data=False, calc_chi2=False, gen_more=None, sample=(X_ if gen_more is None: gen_features = X else: - gen_features = np.tile( - X, - [gen_more] + [1] * (X.ndim - 1) - ) - gen_scaled = np.concatenate([ - model.make_fake(gen_features[i:i+batch_size]).numpy() - for i in range(0, len(gen_features), batch_size) - ], axis=0) + gen_features = np.tile(X, [gen_more] + [1] * (X.ndim - 1)) + gen_scaled = np.concatenate( + [ + model.make_fake(gen_features[i : i + batch_size]).numpy() + for i in range(0, len(gen_features), batch_size) + ], + axis=0, + ) real = unscale(Y) gen = unscale(gen_scaled) gen[gen < 0] = 0 - gen1 = np.where(gen < 1., 0, gen) + gen1 = np.where(gen < 1.0, 0, gen) images = make_metric_plots(real, gen, features={'angle': (X, gen_features)}, calc_chi2=calc_chi2) if calc_chi2: images, chi2 = images @@ -206,15 +206,11 @@ def array_to_img(arr): for part in ['train', 'test']: path = prediction_path / part path.mkdir() - ( - images, images1, img_amplitude, - gen_dataset, chi2 - ) = get_images( - calc_chi2=True, return_raw_data=True, gen_more=10, - sample=( - (X_train, Y_train) if part == 'train' - else (X_test, Y_test) - ) + (images, images1, img_amplitude, gen_dataset, chi2) = get_images( + calc_chi2=True, + return_raw_data=True, + gen_more=10, + sample=((X_train, Y_train) if part == 'train' else (X_test, Y_test)), ) for k, img in images.items(): array_to_img(img).save(str(path / f"{k}.png")) @@ -238,17 +234,28 @@ def array_to_img(arr): else: features_noise = None if args.feature_noise_power is not None: + def features_noise(epoch): - current_power = args.feature_noise_power / (10**(epoch / args.feature_noise_decay)) + current_power = args.feature_noise_power / (10 ** (epoch / args.feature_noise_decay)) with writer_train.as_default(): tf.summary.scalar("features noise power", current_power, epoch) return current_power - train(Y_train, Y_test, model.training_step, model.calculate_losses, args.num_epochs, args.batch_size, - train_writer=writer_train, val_writer=writer_val, - callbacks=[write_hist_summary, save_model, schedule_lr], - features_train=X_train, features_val=X_test, features_noise=features_noise) + train( + Y_train, + Y_test, + model.training_step, + model.calculate_losses, + args.num_epochs, + args.batch_size, + train_writer=writer_train, + val_writer=writer_val, + callbacks=[write_hist_summary, save_model, schedule_lr], + features_train=X_train, + features_val=X_test, + features_noise=features_noise, + ) if __name__ == '__main__': diff --git a/legacy_code/test_script_data_v3.py b/legacy_code/test_script_data_v3.py index 83dd098..b6ea43f 100644 --- a/legacy_code/test_script_data_v3.py +++ b/legacy_code/test_script_data_v3.py @@ -27,8 +27,8 @@ def main(): parser.add_argument('--latent_dim', type=int, default=32, required=False) parser.add_argument('--gpu_num', type=str, required=False) parser.add_argument('--kernel_init', type=str, default='glorot_uniform', required=False) - parser.add_argument('--gp_lambda', type=float, default=10., required=False) - parser.add_argument('--gpdata_lambda', type=float, default=0., required=False) + parser.add_argument('--gp_lambda', type=float, default=10.0, required=False) + parser.add_argument('--gpdata_lambda', type=float, default=0.0, required=False) parser.add_argument('--num_additional_disc_layers', type=int, default=0, required=False) parser.add_argument('--cramer_gan', action='store_true', default=False) parser.add_argument('--features_to_tail', action='store_true', default=True) @@ -40,9 +40,8 @@ def main(): args = parser.parse_args() - assert ( - (args.feature_noise_power is None) == - (args.feature_noise_decay is None) + assert (args.feature_noise_power is None) == ( + args.feature_noise_decay is None ), 'Noise power and decay must be both provided' print("") @@ -81,32 +80,33 @@ def main(): f.write('\n') f.write(f_in.read()) - model = BaselineModel_6x15(kernel_init=args.kernel_init, lr=args.lr, - num_disc_updates=args.num_disc_updates, latent_dim=args.latent_dim, - gp_lambda=args.gp_lambda, gpdata_lambda=args.gpdata_lambda, - num_additional_layers=args.num_additional_disc_layers, - cramer=args.cramer_gan, features_to_tail=args.features_to_tail, - dropout_rate=args.dropout_rate, - stochastic_stepping=args.stochastic_stepping) + model = BaselineModel_6x15( + kernel_init=args.kernel_init, + lr=args.lr, + num_disc_updates=args.num_disc_updates, + latent_dim=args.latent_dim, + gp_lambda=args.gp_lambda, + gpdata_lambda=args.gpdata_lambda, + num_additional_layers=args.num_additional_disc_layers, + cramer=args.cramer_gan, + features_to_tail=args.features_to_tail, + dropout_rate=args.dropout_rate, + stochastic_stepping=args.stochastic_stepping, + ) if args.prediction_only: + def epoch_from_name(name): - epoch, = re.findall(r'\d+', name) + (epoch,) = re.findall(r'\d+', name) return int(epoch) gen_checkpoints = model_path.glob("generator_*.h5") disc_checkpoints = model_path.glob("discriminator_*.h5") - latest_gen_checkpoint = max( - gen_checkpoints, - key=lambda path: epoch_from_name(path.stem) - ) - latest_disc_checkpoint = max( - disc_checkpoints, - key=lambda path: epoch_from_name(path.stem) - ) + latest_gen_checkpoint = max(gen_checkpoints, key=lambda path: epoch_from_name(path.stem)) + latest_disc_checkpoint = max(disc_checkpoints, key=lambda path: epoch_from_name(path.stem)) - assert ( - epoch_from_name(latest_gen_checkpoint.stem) == epoch_from_name(latest_disc_checkpoint.stem) + assert epoch_from_name(latest_gen_checkpoint.stem) == epoch_from_name( + latest_disc_checkpoint.stem ), "Latest disc and gen epochs differ" print(f'Loading generator weights from {str(latest_gen_checkpoint)}') @@ -144,18 +144,18 @@ def get_images(return_raw_data=False, calc_chi2=False, gen_more=None, sample=(X_ if gen_more is None: gen_features = X else: - gen_features = np.tile( - X, - [gen_more] + [1] * (X.ndim - 1) - ) - gen_scaled = np.concatenate([ - model.make_fake(gen_features[i:i+batch_size]).numpy() - for i in range(0, len(gen_features), batch_size) - ], axis=0) + gen_features = np.tile(X, [gen_more] + [1] * (X.ndim - 1)) + gen_scaled = np.concatenate( + [ + model.make_fake(gen_features[i : i + batch_size]).numpy() + for i in range(0, len(gen_features), batch_size) + ], + axis=0, + ) real = unscale(Y) gen = unscale(gen_scaled) gen[gen < 0] = 0 - gen1 = np.where(gen < 1., 0, gen) + gen1 = np.where(gen < 1.0, 0, gen) features = { 'crossing_angle': (X[:, 0], gen_features[:, 0]), @@ -203,6 +203,7 @@ def schedule_lr(step): tf.summary.scalar("generator learning rate", model.gen_opt.lr, step) if args.prediction_only: + def array_to_img(arr): return PIL.Image.fromarray(arr.reshape(arr.shape[1:])) @@ -213,12 +214,11 @@ def array_to_img(arr): for part in ['train', 'test']: path = prediction_path / part path.mkdir() - ( - images, images1, img_amplitude, - gen_dataset, chi2 - ) = get_images( - calc_chi2=True, return_raw_data=True, gen_more=10, - sample=((X_train, Y_train) if part == 'train' else (X_test, Y_test)) + (images, images1, img_amplitude, gen_dataset, chi2) = get_images( + calc_chi2=True, + return_raw_data=True, + gen_more=10, + sample=((X_train, Y_train) if part == 'train' else (X_test, Y_test)), ) for k, img in images.items(): array_to_img(img).save(str(path / f"{k}.png")) @@ -243,17 +243,28 @@ def array_to_img(arr): else: features_noise = None if args.feature_noise_power is not None: + def features_noise(epoch): - current_power = args.feature_noise_power / (10**(epoch / args.feature_noise_decay)) + current_power = args.feature_noise_power / (10 ** (epoch / args.feature_noise_decay)) with writer_train.as_default(): tf.summary.scalar("features noise power", current_power, epoch) return current_power - train(Y_train, Y_test, model.training_step, model.calculate_losses, args.num_epochs, args.batch_size, - train_writer=writer_train, val_writer=writer_val, - callbacks=[write_hist_summary, save_model, schedule_lr], - features_train=X_train, features_val=X_test, features_noise=features_noise) + train( + Y_train, + Y_test, + model.training_step, + model.calculate_losses, + args.num_epochs, + args.batch_size, + train_writer=writer_train, + val_writer=writer_val, + callbacks=[write_hist_summary, save_model, schedule_lr], + features_train=X_train, + features_val=X_test, + features_noise=features_noise, + ) if __name__ == '__main__': diff --git a/legacy_code/test_script_data_v4.py b/legacy_code/test_script_data_v4.py index 6b0ec29..deddad2 100644 --- a/legacy_code/test_script_data_v4.py +++ b/legacy_code/test_script_data_v4.py @@ -27,8 +27,8 @@ def make_parser(): parser.add_argument('--latent_dim', type=int, default=32, required=False) parser.add_argument('--gpu_num', type=str, required=False) parser.add_argument('--kernel_init', type=str, default='glorot_uniform', required=False) - parser.add_argument('--gp_lambda', type=float, default=10., required=False) - parser.add_argument('--gpdata_lambda', type=float, default=0., required=False) + parser.add_argument('--gp_lambda', type=float, default=10.0, required=False) + parser.add_argument('--gpdata_lambda', type=float, default=0.0, required=False) parser.add_argument('--num_additional_disc_layers', type=int, default=0, required=False) parser.add_argument('--cramer_gan', action='store_true', default=False) parser.add_argument('--features_to_tail', action='store_true', default=True) @@ -53,9 +53,8 @@ def print_args(args): def parse_args(): args = make_parser().parse_args() - assert ( - (args.feature_noise_power is None) == - (args.feature_noise_decay is None) + assert (args.feature_noise_power is None) == ( + args.feature_noise_decay is None ), 'Noise power and decay must be both provided' print_args(args) @@ -76,24 +75,18 @@ def write_args(model_path, fname='arguments.txt'): def epoch_from_name(name): - epoch, = re.findall(r'\d+', name) + (epoch,) = re.findall(r'\d+', name) return int(epoch) def load_weights(model, model_path): gen_checkpoints = model_path.glob("generator_*.h5") disc_checkpoints = model_path.glob("discriminator_*.h5") - latest_gen_checkpoint = max( - gen_checkpoints, - key=lambda path: epoch_from_name(path.stem) - ) - latest_disc_checkpoint = max( - disc_checkpoints, - key=lambda path: epoch_from_name(path.stem) - ) + latest_gen_checkpoint = max(gen_checkpoints, key=lambda path: epoch_from_name(path.stem)) + latest_disc_checkpoint = max(disc_checkpoints, key=lambda path: epoch_from_name(path.stem)) - assert ( - epoch_from_name(latest_gen_checkpoint.stem) == epoch_from_name(latest_disc_checkpoint.stem) + assert epoch_from_name(latest_gen_checkpoint.stem) == epoch_from_name( + latest_disc_checkpoint.stem ), "Latest disc and gen epochs differ" print(f'Loading generator weights from {str(latest_gen_checkpoint)}') @@ -104,12 +97,7 @@ def load_weights(model, model_path): return latest_gen_checkpoint, latest_disc_checkpoint -def get_images(model, - sample, - return_raw_data=False, - calc_chi2=False, - gen_more=None, - batch_size=128): +def get_images(model, sample, return_raw_data=False, calc_chi2=False, gen_more=None, batch_size=128): X, Y = sample assert X.ndim == 2 assert X.shape[1] == 4 @@ -117,25 +105,22 @@ def get_images(model, if gen_more is None: gen_features = X else: - gen_features = np.tile( - X, - [gen_more] + [1] * (X.ndim - 1) - ) - gen_scaled = np.concatenate([ - model.make_fake(gen_features[i:i+batch_size]).numpy() - for i in range(0, len(gen_features), batch_size) - ], axis=0) + gen_features = np.tile(X, [gen_more] + [1] * (X.ndim - 1)) + gen_scaled = np.concatenate( + [model.make_fake(gen_features[i : i + batch_size]).numpy() for i in range(0, len(gen_features), batch_size)], + axis=0, + ) real = model.scaler.unscale(Y) gen = model.scaler.unscale(gen_scaled) gen[gen < 0] = 0 - gen1 = np.where(gen < 1., 0, gen) + gen1 = np.where(gen < 1.0, 0, gen) features = { 'crossing_angle': (X[:, 0], gen_features[:, 0]), 'dip_angle': (X[:, 1], gen_features[:, 1]), 'drift_length': (X[:, 2], gen_features[:, 2]), 'time_bin_fraction': (X[:, 2] % 1, gen_features[:, 2] % 1), - 'pad_coord_fraction': (X[:, 3] % 1, gen_features[:, 3] % 1) + 'pad_coord_fraction': (X[:, 3] % 1, gen_features[:, 3] % 1), } images = make_metric_plots(real, gen, features=features, calc_chi2=calc_chi2) @@ -166,10 +151,8 @@ def __init__(self, model, path, save_period): def __call__(self, step): if step % self.save_period == 0: print(f'Saving model on step {step} to {self.path}') - self.model.generator.save( - str(self.path.joinpath("generator_{:05d}.h5".format(step)))) - self.model.discriminator.save( - str(self.path.joinpath("discriminator_{:05d}.h5".format(step)))) + self.model.generator.save(str(self.path.joinpath("generator_{:05d}.h5".format(step)))) + self.model.discriminator.save(str(self.path.joinpath("discriminator_{:05d}.h5".format(step)))) class WriteHistSummaryCallback: @@ -181,9 +164,7 @@ def __init__(self, model, sample, save_period, writer): def __call__(self, step): if step % self.save_period == 0: - images, images1, img_amplitude, chi2 = get_images(self.model, - sample=self.sample, - calc_chi2=True) + images, images1, img_amplitude, chi2 = get_images(self.model, sample=self.sample, calc_chi2=True) with self.writer.as_default(): tf.summary.scalar("chi2", chi2, step) @@ -213,11 +194,9 @@ def array_to_img(arr): return PIL.Image.fromarray(arr.reshape(arr.shape[1:])) path.mkdir() - ( - images, images1, img_amplitude, - gen_dataset, chi2 - ) = get_images(model, sample=sample, - calc_chi2=True, return_raw_data=True, gen_more=10) + (images, images1, img_amplitude, gen_dataset, chi2) = get_images( + model, sample=sample, calc_chi2=True, return_raw_data=True, gen_more=10 + ) for k, img in images.items(): array_to_img(img).save(str(path / f"{k}.png")) @@ -255,13 +234,19 @@ def main(): write_args(model_path) - model = BaselineModel_8x16(kernel_init=args.kernel_init, lr=args.lr, - num_disc_updates=args.num_disc_updates, latent_dim=args.latent_dim, - gp_lambda=args.gp_lambda, gpdata_lambda=args.gpdata_lambda, - num_additional_layers=args.num_additional_disc_layers, - cramer=args.cramer_gan, features_to_tail=args.features_to_tail, - dropout_rate=args.dropout_rate, - stochastic_stepping=args.stochastic_stepping) + model = BaselineModel_8x16( + kernel_init=args.kernel_init, + lr=args.lr, + num_disc_updates=args.num_disc_updates, + latent_dim=args.latent_dim, + gp_lambda=args.gp_lambda, + gpdata_lambda=args.gpdata_lambda, + num_additional_layers=args.num_additional_disc_layers, + cramer=args.cramer_gan, + features_to_tail=args.features_to_tail, + dropout_rate=args.dropout_rate, + stochastic_stepping=args.stochastic_stepping, + ) if args.prediction_only: latest_gen_checkpoint, latest_disc_checkpoint = load_weights(model, model_path) @@ -285,38 +270,42 @@ def main(): for part in ['train', 'test']: evaluate_model( - model, path=prediction_path / part, - sample=( - (X_train, Y_train) if part == 'train' - else (X_test, Y_test) - ), - gen_sample_name=(None if part == 'train' else 'generated.dat') + model, + path=prediction_path / part, + sample=((X_train, Y_train) if part == 'train' else (X_test, Y_test)), + gen_sample_name=(None if part == 'train' else 'generated.dat'), ) else: features_noise = None if args.feature_noise_power is not None: + def features_noise(epoch): - current_power = args.feature_noise_power / (10**(epoch / args.feature_noise_decay)) + current_power = args.feature_noise_power / (10 ** (epoch / args.feature_noise_decay)) with writer_train.as_default(): tf.summary.scalar("features noise power", current_power, epoch) return current_power - save_model = SaveModelCallback( - model=model, path=model_path, save_period=args.save_every - ) + save_model = SaveModelCallback(model=model, path=model_path, save_period=args.save_every) write_hist_summary = WriteHistSummaryCallback( - model, sample=(X_test, Y_test), - save_period=args.save_every, writer=writer_val + model, sample=(X_test, Y_test), save_period=args.save_every, writer=writer_val ) - schedule_lr = ScheduleLRCallback( - model, decay_rate=args.lr_schedule_rate, writer=writer_val + schedule_lr = ScheduleLRCallback(model, decay_rate=args.lr_schedule_rate, writer=writer_val) + train( + Y_train, + Y_test, + model.training_step, + model.calculate_losses, + args.num_epochs, + args.batch_size, + train_writer=writer_train, + val_writer=writer_val, + callbacks=[write_hist_summary, save_model, schedule_lr], + features_train=X_train, + features_val=X_test, + features_noise=features_noise, ) - train(Y_train, Y_test, model.training_step, model.calculate_losses, args.num_epochs, args.batch_size, - train_writer=writer_train, val_writer=writer_val, - callbacks=[write_hist_summary, save_model, schedule_lr], - features_train=X_train, features_val=X_test, features_noise=features_noise) if __name__ == '__main__': diff --git a/metrics/__init__.py b/metrics/__init__.py index 4a1e8be..b44cb5c 100644 --- a/metrics/__init__.py +++ b/metrics/__init__.py @@ -62,14 +62,12 @@ def make_metric_plots(images_real, images_gen, features=None, calc_chi2=False, m pdffile = io.BytesIO() pdf_plots[name] = pdffile if calc_chi2 and (metric_name != "Sum"): - plots[name], chi2_i = make_trend_plot(feature_real, real, - feature_gen, gen, - name, calc_chi2=True, - pdffile=pdffile) + plots[name], chi2_i = make_trend_plot( + feature_real, real, feature_gen, gen, name, calc_chi2=True, pdffile=pdffile + ) chi2 += chi2_i else: - plots[name] = make_trend_plot(feature_real, real, - feature_gen, gen, name, pdffile=pdffile) + plots[name] = make_trend_plot(feature_real, real, feature_gen, gen, name, pdffile=pdffile) except AssertionError as e: print(f"WARNING! Assertion error ({e})") @@ -83,17 +81,13 @@ def make_metric_plots(images_real, images_gen, features=None, calc_chi2=False, m return result -def make_images_for_model(model, - sample, - return_raw_data=False, - calc_chi2=False, - gen_more=None, - batch_size=128, - pdf_outputs=None): +def make_images_for_model( + model, sample, return_raw_data=False, calc_chi2=False, gen_more=None, batch_size=128, pdf_outputs=None +): X, Y = sample assert X.ndim == 2 assert X.shape[1] == 4 - make_pdfs = (pdf_outputs is not None) + make_pdfs = pdf_outputs is not None if make_pdfs: assert isinstance(pdf_outputs, list) assert len(pdf_outputs) == 0 @@ -101,29 +95,25 @@ def make_images_for_model(model, if gen_more is None: gen_features = X else: - gen_features = np.tile( - X, - [gen_more] + [1] * (X.ndim - 1) - ) - gen_scaled = np.concatenate([ - model.make_fake(gen_features[i:i+batch_size]).numpy() - for i in range(0, len(gen_features), batch_size) - ], axis=0) + gen_features = np.tile(X, [gen_more] + [1] * (X.ndim - 1)) + gen_scaled = np.concatenate( + [model.make_fake(gen_features[i : i + batch_size]).numpy() for i in range(0, len(gen_features), batch_size)], + axis=0, + ) real = model.scaler.unscale(Y) gen = model.scaler.unscale(gen_scaled) gen[gen < 0] = 0 - gen1 = np.where(gen < 1., 0, gen) + gen1 = np.where(gen < 1.0, 0, gen) features = { 'crossing_angle': (X[:, 0], gen_features[:, 0]), 'dip_angle': (X[:, 1], gen_features[:, 1]), 'drift_length': (X[:, 2], gen_features[:, 2]), 'time_bin_fraction': (X[:, 2] % 1, gen_features[:, 2] % 1), - 'pad_coord_fraction': (X[:, 3] % 1, gen_features[:, 3] % 1) + 'pad_coord_fraction': (X[:, 3] % 1, gen_features[:, 3] % 1), } - metric_plot_results = make_metric_plots(real, gen, features=features, - calc_chi2=calc_chi2, make_pdfs=make_pdfs) + metric_plot_results = make_metric_plots(real, gen, features=features, calc_chi2=calc_chi2, make_pdfs=make_pdfs) images = metric_plot_results['plots'] if calc_chi2: chi2 = metric_plot_results['chi2'] @@ -140,8 +130,9 @@ def make_images_for_model(model, if make_pdfs: pdffile = io.BytesIO() pdf_outputs.append(pdffile) - img_amplitude = make_histograms(Y.flatten(), gen_scaled.flatten(), 'log10(amplitude + 1)', logy=True, - pdffile=pdffile) + img_amplitude = make_histograms( + Y.flatten(), gen_scaled.flatten(), 'log10(amplitude + 1)', logy=True, pdffile=pdffile + ) pdffile_examples = None pdffile_examples_mask = None @@ -170,11 +161,9 @@ def array_to_img(arr): path.mkdir() pdf_outputs = [] - ( - images, images1, img_amplitude, - gen_dataset, chi2 - ) = make_images_for_model(model, sample=sample, - calc_chi2=True, return_raw_data=True, gen_more=10, pdf_outputs=pdf_outputs) + (images, images1, img_amplitude, gen_dataset, chi2) = make_images_for_model( + model, sample=sample, calc_chi2=True, return_raw_data=True, gen_more=10, pdf_outputs=pdf_outputs + ) images_pdf, images1_pdf, img_amplitude_pdf = pdf_outputs for k, img in images.items(): @@ -222,8 +211,7 @@ def plot_individual_images(real, gen, n=10, pdffile=None): size_y = size_x / real.shape[2] * real.shape[1] * n * 1.2 / 4 fig, axx = plt.subplots(n, 4, figsize=(size_x, size_y)) - axx = [(ax[0], ax[1]) for ax in axx] + \ - [(ax[2], ax[3]) for ax in axx] + axx = [(ax[0], ax[1]) for ax in axx] + [(ax[2], ax[3]) for ax in axx] for ax, img_real, img_fake in zip(axx, real, gen): ax[0].imshow(img_real, aspect='auto') @@ -252,9 +240,9 @@ def plot_images_mask(real, gen, pdffile=None): size_y = size_x / real.shape[2] * real.shape[1] * 2.4 fig, [ax0, ax1] = plt.subplots(2, 1, figsize=(size_x, size_y)) - ax0.imshow((real >= 1.).any(axis=0), aspect='auto') + ax0.imshow((real >= 1.0).any(axis=0), aspect='auto') ax0.set_title("real") - ax1.imshow((gen >= 1.).any(axis=0), aspect='auto') + ax1.imshow((gen >= 1.0).any(axis=0), aspect='auto') ax1.set_title("generated") buf = io.BytesIO() diff --git a/metrics/gaussian_metrics.py b/metrics/gaussian_metrics.py index 718691f..eb29ff5 100644 --- a/metrics/gaussian_metrics.py +++ b/metrics/gaussian_metrics.py @@ -7,15 +7,11 @@ def _gaussian_fit(img): assert (img > 0).any(), '_gaussian_fit: blank image' img_n = img / img.sum() - mu = np.fromfunction( - lambda i, j: (img_n[np.newaxis, ...] * np.stack([i, j])).sum(axis=(1, 2)), - shape=img.shape - ) + mu = np.fromfunction(lambda i, j: (img_n[np.newaxis, ...] * np.stack([i, j])).sum(axis=(1, 2)), shape=img.shape) cov = np.fromfunction( - lambda i, j: ( - (img_n[np.newaxis, ...] * np.stack([i * i, j * i, i * j, j * j])).sum(axis=(1, 2)) - ) - np.stack([mu[0]**2, mu[0]*mu[1], mu[0]*mu[1], mu[1]**2]), - shape=img.shape + lambda i, j: ((img_n[np.newaxis, ...] * np.stack([i * i, j * i, i * j, j * j])).sum(axis=(1, 2))) + - np.stack([mu[0] ** 2, mu[0] * mu[1], mu[0] * mu[1], mu[1] ** 2]), + shape=img.shape, ).reshape(2, 2) return mu, cov @@ -49,14 +45,13 @@ def get_val_metric_v(imgs): imgs_n = imgs / imgs.sum(axis=(1, 2), keepdims=True) mu = np.fromfunction( lambda i, j: (imgs_n[:, np.newaxis, ...] * np.stack([i, j])[np.newaxis, ...]).sum(axis=(2, 3)), - shape=imgs.shape[1:] + shape=imgs.shape[1:], ) cov = np.fromfunction( - lambda i, j: ( - (imgs_n[:, np.newaxis, ...] * np.stack([i * i, j * j, i * j])[np.newaxis, ...]).sum(axis=(2, 3)) - ) - np.stack([mu[:, 0]**2, mu[:, 1]**2, mu[:, 0] * mu[:, 1]]).T, - shape=imgs.shape[1:] + lambda i, j: ((imgs_n[:, np.newaxis, ...] * np.stack([i * i, j * j, i * j])[np.newaxis, ...]).sum(axis=(2, 3))) + - np.stack([mu[:, 0] ** 2, mu[:, 1] ** 2, mu[:, 0] * mu[:, 1]]).T, + shape=imgs.shape[1:], ) return np.concatenate([mu, cov, imgs.sum(axis=(1, 2))[:, np.newaxis]], axis=1) diff --git a/metrics/plotting.py b/metrics/plotting.py index 91f5761..ea51f22 100644 --- a/metrics/plotting.py +++ b/metrics/plotting.py @@ -13,12 +13,13 @@ def _bootstrap_error(data, function, num_bs=100): def _get_stats(arr): class Obj: pass + result = Obj() result.mean = arr.mean() result.width = arr.std() - result.mean_err = result.width / (len(arr) - 1)**0.5 + result.mean_err = result.width / (len(arr) - 1) ** 0.5 result.width_err = _bootstrap_error(arr, np.std) return result @@ -26,11 +27,7 @@ class Obj: def compare_two_dists(d_real, d_gen, label, tag=None, nbins=100): ax = plt.gca() - bins = np.linspace( - min(d_real.min(), d_gen.min()), - max(d_real.max(), d_gen.max()), - nbins + 1 - ) + bins = np.linspace(min(d_real.min(), d_gen.min()), max(d_real.max(), d_gen.max()), nbins + 1) stats_real = _get_stats(d_real) stats_gen = _get_stats(d_gen) @@ -41,14 +38,16 @@ def compare_two_dists(d_real, d_gen, label, tag=None, nbins=100): leg_entry = 'gen' plt.hist(d_real, bins=bins, density=True, label='real') - plt.hist(d_gen, bins=bins, density=True, label=leg_entry, histtype='step', linewidth=2.) - - string = '\n'.join([ - f"real: mean = {stats_real.mean :.4f} +/- {stats_real.mean_err :.4f}", - f"gen: mean = {stats_gen.mean :.4f} +/- {stats_gen .mean_err :.4f}", - f"real: std = {stats_real.width:.4f} +/- {stats_real.width_err:.4f}", - f"gen: std = {stats_gen.width:.4f} +/- {stats_gen .width_err:.4f}", - ]) + plt.hist(d_gen, bins=bins, density=True, label=leg_entry, histtype='step', linewidth=2.0) + + string = '\n'.join( + [ + f"real: mean = {stats_real.mean :.4f} +/- {stats_real.mean_err :.4f}", + f"gen: mean = {stats_gen.mean :.4f} +/- {stats_gen .mean_err :.4f}", + f"real: std = {stats_real.width:.4f} +/- {stats_real.width_err:.4f}", + f"gen: std = {stats_gen.width:.4f} +/- {stats_gen .width_err:.4f}", + ] + ) default_family = rcParams['font.family'] rcParams['font.family'] = 'monospace' ax.add_artist(AnchoredText(string, loc=2)) diff --git a/metrics/trends.py b/metrics/trends.py index 42dab58..3121971 100644 --- a/metrics/trends.py +++ b/metrics/trends.py @@ -16,29 +16,22 @@ def calc_trend(x, y, do_plot=True, bins=100, window_size=20, **kwargs): if isinstance(bins, int): bins = np.linspace(np.min(x), np.max(x), bins + 1) - sel = (x >= bins[0]) + sel = x >= bins[0] x, y = x[sel], y[sel] cats = (x[:, np.newaxis] < bins[np.newaxis, 1:]).argmax(axis=1) def stats(arr): - return ( - arr.mean(), - arr.std() / (len(arr) - 1)**0.5, - arr.std(), - _bootstrap_error(arr, np.std) - ) + return (arr.mean(), arr.std() / (len(arr) - 1) ** 0.5, arr.std(), _bootstrap_error(arr, np.std)) - mean, mean_err, std, std_err, bin_centers = np.array([ - stats( - y[(cats >= left) & (cats < right)] - ) + ((bins[left] + bins[right]) / 2,) for left, right in zip( - range(len(bins) - window_size), - range(window_size, len(bins)) - ) - ]).T + mean, mean_err, std, std_err, bin_centers = np.array( + [ + stats(y[(cats >= left) & (cats < right)]) + ((bins[left] + bins[right]) / 2,) + for left, right in zip(range(len(bins) - window_size), range(window_size, len(bins))) + ] + ).T if do_plot: - mean_p_std_err = (mean_err**2 + std_err**2)**0.5 + mean_p_std_err = (mean_err ** 2 + std_err ** 2) ** 0.5 plt.fill_between(bin_centers, mean - mean_err, mean + mean_err, **kwargs) kwargs['alpha'] *= 0.5 kwargs = {k: v for k, v in kwargs.items() if k != 'label'} @@ -56,11 +49,7 @@ def make_trend_plot(feature_real, real, feature_gen, gen, name, calc_chi2=False, real = real.squeeze() gen = gen.squeeze() - bins = np.linspace( - min(feature_real.min(), feature_gen.min()), - max(feature_real.max(), feature_gen.max()), - 100 - ) + bins = np.linspace(min(feature_real.min(), feature_gen.min()), max(feature_real.max(), feature_gen.max()), 100) fig = plt.figure(figsize=figsize) calc_trend(feature_real, real, bins=bins, label='real', color='blue') @@ -79,32 +68,25 @@ def make_trend_plot(feature_real, real, feature_gen, gen, name, calc_chi2=False, img_data = np.array(img.getdata(), dtype=np.uint8).reshape(1, img.size[1], img.size[0], -1) if calc_chi2: - bins = np.linspace( - min(feature_real.min(), feature_gen.min()), - max(feature_real.max(), feature_gen.max()), - 20 + bins = np.linspace(min(feature_real.min(), feature_gen.min()), max(feature_real.max(), feature_gen.max()), 20) + ((real_mean, real_std), (real_mean_err, real_std_err)) = calc_trend( + feature_real, real, do_plot=False, bins=bins, window_size=1 + ) + ((gen_mean, gen_std), (gen_mean_err, gen_std_err)) = calc_trend( + feature_gen, gen, do_plot=False, bins=bins, window_size=1 ) - ( - (real_mean, real_std), - (real_mean_err, real_std_err) - ) = calc_trend(feature_real, real, do_plot=False, bins=bins, window_size=1) - ( - (gen_mean, gen_std), - (gen_mean_err, gen_std_err) - ) = calc_trend(feature_gen, gen, do_plot=False, bins=bins, window_size=1) gen_upper = gen_mean + gen_std gen_lower = gen_mean - gen_std - gen_err2 = gen_mean_err**2 + gen_std_err**2 + gen_err2 = gen_mean_err ** 2 + gen_std_err ** 2 real_upper = real_mean + real_std real_lower = real_mean - real_std - real_err2 = real_mean_err**2 + real_std_err**2 + real_err2 = real_mean_err ** 2 + real_std_err ** 2 - chi2 = ( - ((gen_upper - real_upper)**2 / (gen_err2 + real_err2)).sum() + - ((gen_lower - real_lower)**2 / (gen_err2 + real_err2)).sum() - ) + chi2 = ((gen_upper - real_upper) ** 2 / (gen_err2 + real_err2)).sum() + ( + (gen_lower - real_lower) ** 2 / (gen_err2 + real_err2) + ).sum() return img_data, chi2 diff --git a/model_export/dump_graph.py b/model_export/dump_graph.py index 8852cfe..4ff3ba3 100644 --- a/model_export/dump_graph.py +++ b/model_export/dump_graph.py @@ -10,20 +10,24 @@ import tf2xla_pb2 -def model_to_graph(model, preprocess, postprocess, input_signature, output_file, test_input=None, - hack_upsampling=False, batch_sizes=(1, 10, 100, 1000, 10000), perf_iterations=5): +def model_to_graph( + model, + preprocess, + postprocess, + input_signature, + output_file, + test_input=None, + hack_upsampling=False, + batch_sizes=(1, 10, 100, 1000, 10000), + perf_iterations=5, +): tf.keras.backend.set_learning_phase(0) @tf.function(input_signature=input_signature) def to_save(x): - return postprocess( - model(preprocess(x)) - ) + return postprocess(model(preprocess(x))) - constant_graph = \ - convert_to_constants.convert_variables_to_constants_v2( - to_save.get_concrete_function() - ) + constant_graph = convert_to_constants.convert_variables_to_constants_v2(to_save.get_concrete_function()) if hack_upsampling: print("Warning: hacking upsampling operations") @@ -40,13 +44,10 @@ def to_save(x): constant_graph.graph.as_graph_def(), [i.op.name for i in constant_graph.inputs], [o.op.name for o in constant_graph.outputs], - tf.float32.as_datatype_enum + tf.float32.as_datatype_enum, ) - tf.io.write_graph( - optimized_graph, - path, filename - ) + tf.io.write_graph(optimized_graph, path, filename) for batch_size in batch_sizes: config = tf2xla_pb2.Config() @@ -66,24 +67,19 @@ def to_save(x): f.write(str(config)) if test_input is not None: - print(to_save( - tf.convert_to_tensor([test_input]) - )) + print(to_save(tf.convert_to_tensor([test_input]))) for batch_size in batch_sizes[::-1]: timings = [] iterations = perf_iterations * max(1, 100 // batch_size) for i in range(iterations): - batched_input = tf.random.normal( - shape=(batch_size, len(test_input)), - dtype='float32' - ) + batched_input = tf.random.normal(shape=(batch_size, len(test_input)), dtype='float32') t0 = perf_counter() to_save(batched_input).numpy() t1 = perf_counter() - timings.append((t1 - t0) * 1000. / batch_size) + timings.append((t1 - t0) * 1000.0 / batch_size) timings = np.array(timings) mean = timings.mean() - err = timings.std() / (len(timings) - 1)**0.5 + err = timings.std() / (len(timings) - 1) ** 0.5 print(f'With batch size = {batch_size}, duration per 1 generation is: {mean} +\\- {err} ms') diff --git a/model_export/tf2xla_pb2.py b/model_export/tf2xla_pb2.py index 8a72a04..5b7bc0e 100644 --- a/model_export/tf2xla_pb2.py +++ b/model_export/tf2xla_pb2.py @@ -6,6 +6,7 @@ from google.protobuf import message as _message from google.protobuf import reflection as _reflection from google.protobuf import symbol_database as _symbol_database + # @@protoc_insertion_point(imports) _sym_db = _symbol_database.Default() @@ -16,298 +17,501 @@ DESCRIPTOR = _descriptor.FileDescriptor( - name='tensorflow/compiler/tf2xla/tf2xla.proto', - package='tensorflow.tf2xla', - syntax='proto3', - serialized_options=b'\n\025org.tensorflow.tf2xlaB\014Tf2XlaProtosP\001\370\001\001', - serialized_pb=b'\n\'tensorflow/compiler/tf2xla/tf2xla.proto\x12\x11tensorflow.tf2xla\x1a,tensorflow/core/framework/tensor_shape.proto\x1a%tensorflow/core/framework/types.proto\"3\n\x08TensorId\x12\x11\n\tnode_name\x18\x01 \x01(\t\x12\x14\n\x0coutput_index\x18\x02 \x01(\x03\"\x8e\x01\n\x04\x46\x65\x65\x64\x12\'\n\x02id\x18\x01 \x01(\x0b\x32\x1b.tensorflow.tf2xla.TensorId\x12+\n\x05shape\x18\x02 \x01(\x0b\x32\x1c.tensorflow.TensorShapeProto\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\"\n\x04type\x18\x04 \x01(\x0e\x32\x14.tensorflow.DataType\"\x8f\x01\n\x05\x46\x65tch\x12\'\n\x02id\x18\x01 \x01(\x0b\x32\x1b.tensorflow.tf2xla.TensorId\x12\x0c\n\x04name\x18\x02 \x01(\t\x12+\n\x05shape\x18\x03 \x01(\x0b\x32\x1c.tensorflow.TensorShapeProto\x12\"\n\x04type\x18\x04 \x01(\x0e\x32\x14.tensorflow.DataType\"\x8e\x01\n\x08Variable\x12\x11\n\tnode_name\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12+\n\x05shape\x18\x03 \x01(\x0b\x32\x1c.tensorflow.TensorShapeProto\x12\"\n\x04type\x18\x04 \x01(\x0e\x32\x14.tensorflow.DataType\x12\x10\n\x08readonly\x18\x05 \x01(\x08\"7\n\x11\x43onversionOptions\x12\"\n\x1a\x63ustom_fake_quant_op_calls\x18\x01 \x01(\x08\"\xc9\x01\n\x06\x43onfig\x12%\n\x04\x66\x65\x65\x64\x18\x01 \x03(\x0b\x32\x17.tensorflow.tf2xla.Feed\x12\'\n\x05\x66\x65tch\x18\x02 \x03(\x0b\x32\x18.tensorflow.tf2xla.Fetch\x12-\n\x08variable\x18\x03 \x03(\x0b\x32\x1b.tensorflow.tf2xla.Variable\x12@\n\x12\x63onversion_options\x18\x04 \x01(\x0b\x32$.tensorflow.tf2xla.ConversionOptionsB*\n\x15org.tensorflow.tf2xlaB\x0cTf2XlaProtosP\x01\xf8\x01\x01\x62\x06proto3' - , - dependencies=[tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2.DESCRIPTOR,tensorflow_dot_core_dot_framework_dot_types__pb2.DESCRIPTOR,]) - - + name='tensorflow/compiler/tf2xla/tf2xla.proto', + package='tensorflow.tf2xla', + syntax='proto3', + serialized_options=b'\n\025org.tensorflow.tf2xlaB\014Tf2XlaProtosP\001\370\001\001', + serialized_pb=b'\n\'tensorflow/compiler/tf2xla/tf2xla.proto\x12\x11tensorflow.tf2xla\x1a,tensorflow/core/framework/tensor_shape.proto\x1a%tensorflow/core/framework/types.proto\"3\n\x08TensorId\x12\x11\n\tnode_name\x18\x01 \x01(\t\x12\x14\n\x0coutput_index\x18\x02 \x01(\x03\"\x8e\x01\n\x04\x46\x65\x65\x64\x12\'\n\x02id\x18\x01 \x01(\x0b\x32\x1b.tensorflow.tf2xla.TensorId\x12+\n\x05shape\x18\x02 \x01(\x0b\x32\x1c.tensorflow.TensorShapeProto\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\"\n\x04type\x18\x04 \x01(\x0e\x32\x14.tensorflow.DataType\"\x8f\x01\n\x05\x46\x65tch\x12\'\n\x02id\x18\x01 \x01(\x0b\x32\x1b.tensorflow.tf2xla.TensorId\x12\x0c\n\x04name\x18\x02 \x01(\t\x12+\n\x05shape\x18\x03 \x01(\x0b\x32\x1c.tensorflow.TensorShapeProto\x12\"\n\x04type\x18\x04 \x01(\x0e\x32\x14.tensorflow.DataType\"\x8e\x01\n\x08Variable\x12\x11\n\tnode_name\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12+\n\x05shape\x18\x03 \x01(\x0b\x32\x1c.tensorflow.TensorShapeProto\x12\"\n\x04type\x18\x04 \x01(\x0e\x32\x14.tensorflow.DataType\x12\x10\n\x08readonly\x18\x05 \x01(\x08\"7\n\x11\x43onversionOptions\x12\"\n\x1a\x63ustom_fake_quant_op_calls\x18\x01 \x01(\x08\"\xc9\x01\n\x06\x43onfig\x12%\n\x04\x66\x65\x65\x64\x18\x01 \x03(\x0b\x32\x17.tensorflow.tf2xla.Feed\x12\'\n\x05\x66\x65tch\x18\x02 \x03(\x0b\x32\x18.tensorflow.tf2xla.Fetch\x12-\n\x08variable\x18\x03 \x03(\x0b\x32\x1b.tensorflow.tf2xla.Variable\x12@\n\x12\x63onversion_options\x18\x04 \x01(\x0b\x32$.tensorflow.tf2xla.ConversionOptionsB*\n\x15org.tensorflow.tf2xlaB\x0cTf2XlaProtosP\x01\xf8\x01\x01\x62\x06proto3', + dependencies=[ + tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2.DESCRIPTOR, + tensorflow_dot_core_dot_framework_dot_types__pb2.DESCRIPTOR, + ], +) _TENSORID = _descriptor.Descriptor( - name='TensorId', - full_name='tensorflow.tf2xla.TensorId', - filename=None, - file=DESCRIPTOR, - containing_type=None, - fields=[ - _descriptor.FieldDescriptor( - name='node_name', full_name='tensorflow.tf2xla.TensorId.node_name', index=0, - number=1, type=9, cpp_type=9, label=1, - has_default_value=False, default_value=b"".decode('utf-8'), - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='output_index', full_name='tensorflow.tf2xla.TensorId.output_index', index=1, - number=2, type=3, cpp_type=2, label=1, - has_default_value=False, default_value=0, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - ], - extensions=[ - ], - nested_types=[], - enum_types=[ - ], - serialized_options=None, - is_extendable=False, - syntax='proto3', - extension_ranges=[], - oneofs=[ - ], - serialized_start=147, - serialized_end=198, + name='TensorId', + full_name='tensorflow.tf2xla.TensorId', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='node_name', + full_name='tensorflow.tf2xla.TensorId.node_name', + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode('utf-8'), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='output_index', + full_name='tensorflow.tf2xla.TensorId.output_index', + index=1, + number=2, + type=3, + cpp_type=2, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + serialized_options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[], + serialized_start=147, + serialized_end=198, ) _FEED = _descriptor.Descriptor( - name='Feed', - full_name='tensorflow.tf2xla.Feed', - filename=None, - file=DESCRIPTOR, - containing_type=None, - fields=[ - _descriptor.FieldDescriptor( - name='id', full_name='tensorflow.tf2xla.Feed.id', index=0, - number=1, type=11, cpp_type=10, label=1, - has_default_value=False, default_value=None, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='shape', full_name='tensorflow.tf2xla.Feed.shape', index=1, - number=2, type=11, cpp_type=10, label=1, - has_default_value=False, default_value=None, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='name', full_name='tensorflow.tf2xla.Feed.name', index=2, - number=3, type=9, cpp_type=9, label=1, - has_default_value=False, default_value=b"".decode('utf-8'), - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='type', full_name='tensorflow.tf2xla.Feed.type', index=3, - number=4, type=14, cpp_type=8, label=1, - has_default_value=False, default_value=0, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - ], - extensions=[ - ], - nested_types=[], - enum_types=[ - ], - serialized_options=None, - is_extendable=False, - syntax='proto3', - extension_ranges=[], - oneofs=[ - ], - serialized_start=201, - serialized_end=343, + name='Feed', + full_name='tensorflow.tf2xla.Feed', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='id', + full_name='tensorflow.tf2xla.Feed.id', + index=0, + number=1, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='shape', + full_name='tensorflow.tf2xla.Feed.shape', + index=1, + number=2, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='name', + full_name='tensorflow.tf2xla.Feed.name', + index=2, + number=3, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode('utf-8'), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='type', + full_name='tensorflow.tf2xla.Feed.type', + index=3, + number=4, + type=14, + cpp_type=8, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + serialized_options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[], + serialized_start=201, + serialized_end=343, ) _FETCH = _descriptor.Descriptor( - name='Fetch', - full_name='tensorflow.tf2xla.Fetch', - filename=None, - file=DESCRIPTOR, - containing_type=None, - fields=[ - _descriptor.FieldDescriptor( - name='id', full_name='tensorflow.tf2xla.Fetch.id', index=0, - number=1, type=11, cpp_type=10, label=1, - has_default_value=False, default_value=None, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='name', full_name='tensorflow.tf2xla.Fetch.name', index=1, - number=2, type=9, cpp_type=9, label=1, - has_default_value=False, default_value=b"".decode('utf-8'), - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='shape', full_name='tensorflow.tf2xla.Fetch.shape', index=2, - number=3, type=11, cpp_type=10, label=1, - has_default_value=False, default_value=None, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='type', full_name='tensorflow.tf2xla.Fetch.type', index=3, - number=4, type=14, cpp_type=8, label=1, - has_default_value=False, default_value=0, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - ], - extensions=[ - ], - nested_types=[], - enum_types=[ - ], - serialized_options=None, - is_extendable=False, - syntax='proto3', - extension_ranges=[], - oneofs=[ - ], - serialized_start=346, - serialized_end=489, + name='Fetch', + full_name='tensorflow.tf2xla.Fetch', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='id', + full_name='tensorflow.tf2xla.Fetch.id', + index=0, + number=1, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='name', + full_name='tensorflow.tf2xla.Fetch.name', + index=1, + number=2, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode('utf-8'), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='shape', + full_name='tensorflow.tf2xla.Fetch.shape', + index=2, + number=3, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='type', + full_name='tensorflow.tf2xla.Fetch.type', + index=3, + number=4, + type=14, + cpp_type=8, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + serialized_options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[], + serialized_start=346, + serialized_end=489, ) _VARIABLE = _descriptor.Descriptor( - name='Variable', - full_name='tensorflow.tf2xla.Variable', - filename=None, - file=DESCRIPTOR, - containing_type=None, - fields=[ - _descriptor.FieldDescriptor( - name='node_name', full_name='tensorflow.tf2xla.Variable.node_name', index=0, - number=1, type=9, cpp_type=9, label=1, - has_default_value=False, default_value=b"".decode('utf-8'), - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='name', full_name='tensorflow.tf2xla.Variable.name', index=1, - number=2, type=9, cpp_type=9, label=1, - has_default_value=False, default_value=b"".decode('utf-8'), - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='shape', full_name='tensorflow.tf2xla.Variable.shape', index=2, - number=3, type=11, cpp_type=10, label=1, - has_default_value=False, default_value=None, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='type', full_name='tensorflow.tf2xla.Variable.type', index=3, - number=4, type=14, cpp_type=8, label=1, - has_default_value=False, default_value=0, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='readonly', full_name='tensorflow.tf2xla.Variable.readonly', index=4, - number=5, type=8, cpp_type=7, label=1, - has_default_value=False, default_value=False, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - ], - extensions=[ - ], - nested_types=[], - enum_types=[ - ], - serialized_options=None, - is_extendable=False, - syntax='proto3', - extension_ranges=[], - oneofs=[ - ], - serialized_start=492, - serialized_end=634, + name='Variable', + full_name='tensorflow.tf2xla.Variable', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='node_name', + full_name='tensorflow.tf2xla.Variable.node_name', + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode('utf-8'), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='name', + full_name='tensorflow.tf2xla.Variable.name', + index=1, + number=2, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode('utf-8'), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='shape', + full_name='tensorflow.tf2xla.Variable.shape', + index=2, + number=3, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='type', + full_name='tensorflow.tf2xla.Variable.type', + index=3, + number=4, + type=14, + cpp_type=8, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='readonly', + full_name='tensorflow.tf2xla.Variable.readonly', + index=4, + number=5, + type=8, + cpp_type=7, + label=1, + has_default_value=False, + default_value=False, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + serialized_options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[], + serialized_start=492, + serialized_end=634, ) _CONVERSIONOPTIONS = _descriptor.Descriptor( - name='ConversionOptions', - full_name='tensorflow.tf2xla.ConversionOptions', - filename=None, - file=DESCRIPTOR, - containing_type=None, - fields=[ - _descriptor.FieldDescriptor( - name='custom_fake_quant_op_calls', full_name='tensorflow.tf2xla.ConversionOptions.custom_fake_quant_op_calls', index=0, - number=1, type=8, cpp_type=7, label=1, - has_default_value=False, default_value=False, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - ], - extensions=[ - ], - nested_types=[], - enum_types=[ - ], - serialized_options=None, - is_extendable=False, - syntax='proto3', - extension_ranges=[], - oneofs=[ - ], - serialized_start=636, - serialized_end=691, + name='ConversionOptions', + full_name='tensorflow.tf2xla.ConversionOptions', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='custom_fake_quant_op_calls', + full_name='tensorflow.tf2xla.ConversionOptions.custom_fake_quant_op_calls', + index=0, + number=1, + type=8, + cpp_type=7, + label=1, + has_default_value=False, + default_value=False, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + serialized_options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[], + serialized_start=636, + serialized_end=691, ) _CONFIG = _descriptor.Descriptor( - name='Config', - full_name='tensorflow.tf2xla.Config', - filename=None, - file=DESCRIPTOR, - containing_type=None, - fields=[ - _descriptor.FieldDescriptor( - name='feed', full_name='tensorflow.tf2xla.Config.feed', index=0, - number=1, type=11, cpp_type=10, label=3, - has_default_value=False, default_value=[], - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='fetch', full_name='tensorflow.tf2xla.Config.fetch', index=1, - number=2, type=11, cpp_type=10, label=3, - has_default_value=False, default_value=[], - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='variable', full_name='tensorflow.tf2xla.Config.variable', index=2, - number=3, type=11, cpp_type=10, label=3, - has_default_value=False, default_value=[], - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - _descriptor.FieldDescriptor( - name='conversion_options', full_name='tensorflow.tf2xla.Config.conversion_options', index=3, - number=4, type=11, cpp_type=10, label=1, - has_default_value=False, default_value=None, - message_type=None, enum_type=None, containing_type=None, - is_extension=False, extension_scope=None, - serialized_options=None, file=DESCRIPTOR), - ], - extensions=[ - ], - nested_types=[], - enum_types=[ - ], - serialized_options=None, - is_extendable=False, - syntax='proto3', - extension_ranges=[], - oneofs=[ - ], - serialized_start=694, - serialized_end=895, + name='Config', + full_name='tensorflow.tf2xla.Config', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='feed', + full_name='tensorflow.tf2xla.Config.feed', + index=0, + number=1, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='fetch', + full_name='tensorflow.tf2xla.Config.fetch', + index=1, + number=2, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='variable', + full_name='tensorflow.tf2xla.Config.variable', + index=2, + number=3, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + _descriptor.FieldDescriptor( + name='conversion_options', + full_name='tensorflow.tf2xla.Config.conversion_options', + index=3, + number=4, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + serialized_options=None, + is_extendable=False, + syntax='proto3', + extension_ranges=[], + oneofs=[], + serialized_start=694, + serialized_end=895, ) _FEED.fields_by_name['id'].message_type = _TENSORID @@ -316,7 +520,9 @@ _FETCH.fields_by_name['id'].message_type = _TENSORID _FETCH.fields_by_name['shape'].message_type = tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2._TENSORSHAPEPROTO _FETCH.fields_by_name['type'].enum_type = tensorflow_dot_core_dot_framework_dot_types__pb2._DATATYPE -_VARIABLE.fields_by_name['shape'].message_type = tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2._TENSORSHAPEPROTO +_VARIABLE.fields_by_name[ + 'shape' +].message_type = tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2._TENSORSHAPEPROTO _VARIABLE.fields_by_name['type'].enum_type = tensorflow_dot_core_dot_framework_dot_types__pb2._DATATYPE _CONFIG.fields_by_name['feed'].message_type = _FEED _CONFIG.fields_by_name['fetch'].message_type = _FETCH @@ -330,46 +536,70 @@ DESCRIPTOR.message_types_by_name['Config'] = _CONFIG _sym_db.RegisterFileDescriptor(DESCRIPTOR) -TensorId = _reflection.GeneratedProtocolMessageType('TensorId', (_message.Message,), { - 'DESCRIPTOR' : _TENSORID, - '__module__' : 'tensorflow.compiler.tf2xla.tf2xla_pb2' - # @@protoc_insertion_point(class_scope:tensorflow.tf2xla.TensorId) - }) +TensorId = _reflection.GeneratedProtocolMessageType( + 'TensorId', + (_message.Message,), + { + 'DESCRIPTOR': _TENSORID, + '__module__': 'tensorflow.compiler.tf2xla.tf2xla_pb2' + # @@protoc_insertion_point(class_scope:tensorflow.tf2xla.TensorId) + }, +) _sym_db.RegisterMessage(TensorId) -Feed = _reflection.GeneratedProtocolMessageType('Feed', (_message.Message,), { - 'DESCRIPTOR' : _FEED, - '__module__' : 'tensorflow.compiler.tf2xla.tf2xla_pb2' - # @@protoc_insertion_point(class_scope:tensorflow.tf2xla.Feed) - }) +Feed = _reflection.GeneratedProtocolMessageType( + 'Feed', + (_message.Message,), + { + 'DESCRIPTOR': _FEED, + '__module__': 'tensorflow.compiler.tf2xla.tf2xla_pb2' + # @@protoc_insertion_point(class_scope:tensorflow.tf2xla.Feed) + }, +) _sym_db.RegisterMessage(Feed) -Fetch = _reflection.GeneratedProtocolMessageType('Fetch', (_message.Message,), { - 'DESCRIPTOR' : _FETCH, - '__module__' : 'tensorflow.compiler.tf2xla.tf2xla_pb2' - # @@protoc_insertion_point(class_scope:tensorflow.tf2xla.Fetch) - }) +Fetch = _reflection.GeneratedProtocolMessageType( + 'Fetch', + (_message.Message,), + { + 'DESCRIPTOR': _FETCH, + '__module__': 'tensorflow.compiler.tf2xla.tf2xla_pb2' + # @@protoc_insertion_point(class_scope:tensorflow.tf2xla.Fetch) + }, +) _sym_db.RegisterMessage(Fetch) -Variable = _reflection.GeneratedProtocolMessageType('Variable', (_message.Message,), { - 'DESCRIPTOR' : _VARIABLE, - '__module__' : 'tensorflow.compiler.tf2xla.tf2xla_pb2' - # @@protoc_insertion_point(class_scope:tensorflow.tf2xla.Variable) - }) +Variable = _reflection.GeneratedProtocolMessageType( + 'Variable', + (_message.Message,), + { + 'DESCRIPTOR': _VARIABLE, + '__module__': 'tensorflow.compiler.tf2xla.tf2xla_pb2' + # @@protoc_insertion_point(class_scope:tensorflow.tf2xla.Variable) + }, +) _sym_db.RegisterMessage(Variable) -ConversionOptions = _reflection.GeneratedProtocolMessageType('ConversionOptions', (_message.Message,), { - 'DESCRIPTOR' : _CONVERSIONOPTIONS, - '__module__' : 'tensorflow.compiler.tf2xla.tf2xla_pb2' - # @@protoc_insertion_point(class_scope:tensorflow.tf2xla.ConversionOptions) - }) +ConversionOptions = _reflection.GeneratedProtocolMessageType( + 'ConversionOptions', + (_message.Message,), + { + 'DESCRIPTOR': _CONVERSIONOPTIONS, + '__module__': 'tensorflow.compiler.tf2xla.tf2xla_pb2' + # @@protoc_insertion_point(class_scope:tensorflow.tf2xla.ConversionOptions) + }, +) _sym_db.RegisterMessage(ConversionOptions) -Config = _reflection.GeneratedProtocolMessageType('Config', (_message.Message,), { - 'DESCRIPTOR' : _CONFIG, - '__module__' : 'tensorflow.compiler.tf2xla.tf2xla_pb2' - # @@protoc_insertion_point(class_scope:tensorflow.tf2xla.Config) - }) +Config = _reflection.GeneratedProtocolMessageType( + 'Config', + (_message.Message,), + { + 'DESCRIPTOR': _CONFIG, + '__module__': 'tensorflow.compiler.tf2xla.tf2xla_pb2' + # @@protoc_insertion_point(class_scope:tensorflow.tf2xla.Config) + }, +) _sym_db.RegisterMessage(Config) diff --git a/models/callbacks.py b/models/callbacks.py index 66b5b3b..3574b33 100644 --- a/models/callbacks.py +++ b/models/callbacks.py @@ -12,10 +12,8 @@ def __init__(self, model, path, save_period): def __call__(self, step): if step % self.save_period == 0: print(f'Saving model on step {step} to {self.path}') - self.model.generator.save( - str(self.path.joinpath("generator_{:05d}.h5".format(step)))) - self.model.discriminator.save( - str(self.path.joinpath("discriminator_{:05d}.h5".format(step)))) + self.model.generator.save(str(self.path.joinpath("generator_{:05d}.h5".format(step)))) + self.model.discriminator.save(str(self.path.joinpath("discriminator_{:05d}.h5".format(step)))) class WriteHistSummaryCallback: @@ -27,9 +25,7 @@ def __init__(self, model, sample, save_period, writer): def __call__(self, step): if step % self.save_period == 0: - images, images1, img_amplitude, chi2 = make_images_for_model(self.model, - sample=self.sample, - calc_chi2=True) + images, images1, img_amplitude, chi2 = make_images_for_model(self.model, sample=self.sample, calc_chi2=True) with self.writer.as_default(): tf.summary.scalar("chi2", chi2, step) @@ -60,5 +56,6 @@ def get_scheduler(lr, lr_decay): return eval(lr_decay) def schedule_lr(step): - return lr * lr_decay**step + return lr * lr_decay ** step + return schedule_lr diff --git a/models/model_v4.py b/models/model_v4.py index 788850d..850061b 100644 --- a/models/model_v4.py +++ b/models/model_v4.py @@ -14,9 +14,7 @@ def preprocess_features(features): # drift_length [35, 290] # pad_coordinate [40-something, 40-something] bin_fractions = features[:, -2:] % 1 - features = ( - features[:, :3] - tf.constant([[0., 0., 162.5]]) - ) / tf.constant([[20., 60., 127.5]]) + features = (features[:, :3] - tf.constant([[0.0, 0.0, 162.5]])) / tf.constant([[20.0, 60.0, 127.5]]) return tf.concat([features, bin_fractions], axis=-1) @@ -33,10 +31,10 @@ def gen_loss(d_real, d_fake): def disc_loss_cramer(d_real, d_fake, d_fake_2): return -tf.reduce_mean( - tf.norm(d_real - d_fake, axis=-1) + - tf.norm(d_fake_2, axis=-1) - - tf.norm(d_fake - d_fake_2, axis=-1) - - tf.norm(d_real, axis=-1) + tf.norm(d_real - d_fake, axis=-1) + + tf.norm(d_fake_2, axis=-1) + - tf.norm(d_fake - d_fake_2, axis=-1) + - tf.norm(d_real, axis=-1) ) @@ -49,17 +47,11 @@ def logloss(x): def disc_loss_js(d_real, d_fake): - return tf.reduce_sum( - logloss(d_real) - ) + tf.reduce_sum( - logloss(-d_fake) - ) / (len(d_real) + len(d_fake)) + return tf.reduce_sum(logloss(d_real)) + tf.reduce_sum(logloss(-d_fake)) / (len(d_real) + len(d_fake)) def gen_loss_js(d_real, d_fake): - return tf.reduce_mean( - logloss(d_fake) - ) + return tf.reduce_mean(logloss(d_fake)) class Model_v4: @@ -82,10 +74,12 @@ def __init__(self, config): self.latent_dim = config['latent_dim'] architecture_descr = config['architecture'] - self.generator = nn.build_architecture(architecture_descr['generator'], - custom_objects_code=config.get('custom_objects', None)) - self.discriminator = nn.build_architecture(architecture_descr['discriminator'], - custom_objects_code=config.get('custom_objects', None)) + self.generator = nn.build_architecture( + architecture_descr['generator'], custom_objects_code=config.get('custom_objects', None) + ) + self.discriminator = nn.build_architecture( + architecture_descr['discriminator'], custom_objects_code=config.get('custom_objects', None) + ) self.step_counter = tf.Variable(0, dtype='int32', trainable=False) @@ -126,28 +120,24 @@ def _load_weights(self, checkpoint, gen_or_disc): if 'optimizer_weights' in model_file: print('Also recovering the optimizer state') - opt_weight_values = hdf5_format.load_optimizer_weights_from_hdf5_group( - model_file - ) + opt_weight_values = hdf5_format.load_optimizer_weights_from_hdf5_group(model_file) network.optimizer.set_weights(opt_weight_values) @tf.function def make_fake(self, features): size = tf.shape(features)[0] latent_input = tf.random.normal(shape=(size, self.latent_dim), dtype='float32') - return self.generator( - tf.concat([_f(features), latent_input], axis=-1) - ) + return self.generator(tf.concat([_f(features), latent_input], axis=-1)) def gradient_penalty(self, features, real, fake): - alpha = tf.random.uniform(shape=[len(real), ] + [1] * (len(real.shape) - 1)) + alpha = tf.random.uniform(shape=[len(real)] + [1] * (len(real.shape) - 1)) interpolates = alpha * real + (1 - alpha) * fake with tf.GradientTape() as t: t.watch(interpolates) d_int = self.discriminator([_f(features), interpolates]) 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) + return tf.reduce_mean(tf.maximum(tf.norm(grads, axis=-1) - 1, 0) ** 2) def gradient_penalty_on_data(self, features, real): with tf.GradientTape() as t: @@ -155,7 +145,7 @@ def gradient_penalty_on_data(self, features, real): d_real = self.discriminator([_f(features), real]) grads = tf.reshape(t.gradient(d_real, real), [len(real), -1]) - return tf.reduce_mean(tf.reduce_sum(grads**2, axis=-1)) + return tf.reduce_mean(tf.reduce_sum(grads ** 2, axis=-1)) @tf.function def calculate_losses(self, feature_batch, target_batch): @@ -175,19 +165,9 @@ def calculate_losses(self, feature_batch, target_batch): d_loss = disc_loss_cramer(d_real, d_fake, d_fake_2) if self.gp_lambda > 0: - d_loss = ( - d_loss + - self.gradient_penalty( - feature_batch, target_batch, fake - ) * self.gp_lambda - ) + d_loss = d_loss + self.gradient_penalty(feature_batch, target_batch, fake) * self.gp_lambda if self.gpdata_lambda > 0: - d_loss = ( - d_loss + - self.gradient_penalty_on_data( - feature_batch, target_batch - ) * self.gpdata_lambda - ) + d_loss = d_loss + self.gradient_penalty_on_data(feature_batch, target_batch) * self.gpdata_lambda if not self.cramer: if self.js: g_loss = gen_loss_js(d_real, d_fake) @@ -223,10 +203,7 @@ def gen_step(self, feature_batch, target_batch): @tf.function def training_step(self, feature_batch, target_batch): if self.stochastic_stepping: - if tf.random.uniform( - shape=[], dtype='int32', - maxval=self.num_disc_updates + 1 - ) == self.num_disc_updates: + if tf.random.uniform(shape=[], dtype='int32', maxval=self.num_disc_updates + 1) == self.num_disc_updates: result = self.gen_step(feature_batch, target_batch) else: result = self.disc_step(feature_batch, target_batch) diff --git a/models/nn.py b/models/nn.py index c55c391..3f692ef 100644 --- a/models/nn.py +++ b/models/nn.py @@ -12,9 +12,9 @@ def get_activation(activation): return activation -def fully_connected_block(units, activations, - kernel_init='glorot_uniform', input_shape=None, - output_shape=None, dropouts=None, name=None): +def fully_connected_block( + units, activations, kernel_init='glorot_uniform', input_shape=None, output_shape=None, dropouts=None, name=None +): assert len(units) == len(activations) if dropouts: assert len(dropouts) == len(units) @@ -42,9 +42,16 @@ def fully_connected_block(units, activations, return tf.keras.Sequential(layers, **args) -def fully_connected_residual_block(units, activations, input_shape, - kernel_init='glorot_uniform', batchnorm=True, - output_shape=None, dropouts=None, name=None): +def fully_connected_residual_block( + units, + activations, + input_shape, + kernel_init='glorot_uniform', + batchnorm=True, + output_shape=None, + dropouts=None, + name=None, +): assert isinstance(units, int) if dropouts: assert len(dropouts) == len(activations) @@ -65,8 +72,7 @@ def single_block(xx, units, activation, kernel_init, batchnorm, dropout): input_tensor = tf.keras.Input(shape=input_shape) xx = input_tensor for i, (act, dropout) in enumerate(zip(activations, dropouts)): - args = dict(units=units, activation=act, kernel_init=kernel_init, - batchnorm=batchnorm, dropout=dropout) + args = dict(units=units, activation=act, kernel_init=kernel_init, batchnorm=batchnorm, dropout=dropout) if len(xx.shape) == 2 and xx.shape[1] == units: xx = xx + single_block(xx, **args) else: @@ -82,8 +88,7 @@ def single_block(xx, units, activation, kernel_init, batchnorm, dropout): return tf.keras.Model(**args) -def concat_block(input1_shape, input2_shape, reshape_input1=None, - reshape_input2=None, axis=-1, name=None): +def concat_block(input1_shape, input2_shape, reshape_input1=None, reshape_input2=None, axis=-1, name=None): in1 = tf.keras.Input(shape=input1_shape) in2 = tf.keras.Input(shape=input2_shape) concat1, concat2 = in1, in2 @@ -98,9 +103,18 @@ def concat_block(input1_shape, input2_shape, reshape_input1=None, return tf.keras.Model(**args) -def conv_block(filters, kernel_sizes, paddings, activations, poolings, - kernel_init='glorot_uniform', input_shape=None, output_shape=None, - dropouts=None, name=None): +def conv_block( + filters, + kernel_sizes, + paddings, + activations, + poolings, + kernel_init='glorot_uniform', + input_shape=None, + output_shape=None, + dropouts=None, + name=None, +): assert len(filters) == len(kernel_sizes) == len(paddings) == len(activations) == len(poolings) if dropouts: assert len(dropouts) == len(filters) @@ -108,10 +122,8 @@ def conv_block(filters, kernel_sizes, paddings, activations, poolings, activations = [get_activation(a) for a in activations] layers = [] - for i, (nfilt, ksize, padding, act, pool) in enumerate(zip(filters, kernel_sizes, paddings, - activations, poolings)): - args = dict(filters=nfilt, kernel_size=ksize, - padding=padding, activation=act, kernel_initializer=kernel_init) + for i, (nfilt, ksize, padding, act, pool) in enumerate(zip(filters, kernel_sizes, paddings, activations, poolings)): + args = dict(filters=nfilt, kernel_size=ksize, padding=padding, activation=act, kernel_initializer=kernel_init) if i == 0 and input_shape: args['input_shape'] = input_shape @@ -133,8 +145,7 @@ def conv_block(filters, kernel_sizes, paddings, activations, poolings, return tf.keras.Sequential(layers, **args) -def vector_img_connect_block(vector_shape, img_shape, block, - vector_bypass=False, concat_outputs=True, name=None): +def vector_img_connect_block(vector_shape, img_shape, block, vector_bypass=False, concat_outputs=True, name=None): vector_shape = tuple(vector_shape) img_shape = tuple(img_shape) @@ -148,10 +159,7 @@ def vector_img_connect_block(vector_shape, img_shape, block, if len(img_shape) == 2: block_input = tf.keras.layers.Reshape(img_shape + (1,))(block_input) if not vector_bypass: - reshaped_vec = tf.tile( - tf.keras.layers.Reshape((1, 1) + vector_shape)(input_vec), - (1, *img_shape[:2], 1) - ) + reshaped_vec = tf.tile(tf.keras.layers.Reshape((1, 1) + vector_shape)(input_vec), (1, *img_shape[:2], 1)) block_input = tf.keras.layers.Concatenate(axis=-1)([block_input, reshaped_vec]) block_output = block(block_input) @@ -160,10 +168,7 @@ def vector_img_connect_block(vector_shape, img_shape, block, if concat_outputs: outputs = tf.keras.layers.Concatenate(axis=-1)(outputs) - args = dict( - inputs=[input_vec, input_img], - outputs=outputs, - ) + args = dict(inputs=[input_vec, input_img], outputs=outputs,) if name: args['name'] = name @@ -185,7 +190,7 @@ def build_block(block_type, arguments): elif block_type == 'fully_connected_residual': block = fully_connected_residual_block(**arguments) else: - raise(NotImplementedError(block_type)) + raise (NotImplementedError(block_type)) return block @@ -196,21 +201,14 @@ def build_architecture(block_descriptions, name=None, custom_objects_code=None): print(custom_objects_code) exec(custom_objects_code, globals(), custom_objects) - blocks = [build_block(**descr) - for descr in block_descriptions] + blocks = [build_block(**descr) for descr in block_descriptions] - inputs = [ - tf.keras.Input(shape=i.shape[1:]) - for i in blocks[0].inputs - ] + inputs = [tf.keras.Input(shape=i.shape[1:]) for i in blocks[0].inputs] outputs = inputs for block in blocks: outputs = block(outputs) - args = dict( - inputs=inputs, - outputs=outputs - ) + args = dict(inputs=inputs, outputs=outputs) if name: args['name'] = name return tf.keras.Model(**args) diff --git a/models/scalers.py b/models/scalers.py index 98cca0f..bf91442 100644 --- a/models/scalers.py +++ b/models/scalers.py @@ -26,22 +26,19 @@ def __init__(self, shape=(8, 16)): def scale(self, x): result = gaussian_fit(x) result[:, -1] = np.log1p(result[:, -1]) - result[:, 4] /= (result[:, 2] * result[:, 3]) + result[:, 4] /= result[:, 2] * result[:, 3] return result def unscale(self, x): m0, m1, D00, D11, D01, logA = x.T D00 = np.clip(D00, 0.05, None) D11 = np.clip(D11, 0.05, None) - D01 = np.clip(D01, -1., 1.) + D01 = np.clip(D01, -1.0, 1.0) D01 *= D00 * D11 A = np.expm1(logA) - cov = np.stack([ - np.stack([D00, D01], axis=1), - np.stack([D01, D11], axis=1) - ], axis=2) # N x 2 x 2 + cov = np.stack([np.stack([D00, D01], axis=1), np.stack([D01, D11], axis=1)], axis=2) # N x 2 x 2 invcov = np.linalg.inv(cov) mu = np.stack([m0, m1], axis=1) diff --git a/models/training.py b/models/training.py index 59c59c5..8b11ce1 100644 --- a/models/training.py +++ b/models/training.py @@ -3,9 +3,21 @@ 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=[], features_train=None, features_val=None, - features_noise=None, first_epoch=0): +def train( + data_train, + data_val, + train_step_fn, + loss_eval_fn, + num_epochs, + batch_size, + train_writer=None, + val_writer=None, + callbacks=[], + features_train=None, + features_val=None, + features_noise=None, + first_epoch=0, +): if not ((features_train is None) or (features_val is None)): assert features_train is not None, 'train: features should be provided for both train and val' assert features_val is not None, 'train: features should be provided for both train and val' @@ -22,15 +34,13 @@ def train(data_train, data_val, train_step_fn, loss_eval_fn, num_epochs, batch_s noise_power = features_noise(i_epoch) for i_sample in trange(0, len(data_train), batch_size): - batch = data_train[shuffle_ids][i_sample:i_sample + batch_size] + batch = data_train[shuffle_ids][i_sample : i_sample + batch_size] if features_train is not None: - feature_batch = features_train[shuffle_ids][i_sample:i_sample + batch_size] + feature_batch = features_train[shuffle_ids][i_sample : i_sample + batch_size] if noise_power is not None: feature_batch = ( - feature_batch + - np.random.normal( - size=feature_batch.shape - ).astype(feature_batch.dtype) * noise_power + feature_batch + + np.random.normal(size=feature_batch.shape).astype(feature_batch.dtype) * noise_power ) if features_train is None: @@ -45,12 +55,12 @@ def train(data_train, data_val, train_step_fn, loss_eval_fn, num_epochs, batch_s losses_val = {} for i_sample in trange(0, len(data_val), batch_size): - batch = data_val[i_sample:i_sample + batch_size] + batch = data_val[i_sample : i_sample + batch_size] if features_train is None: losses_val_batch = {k: l.numpy() for k, l in loss_eval_fn(batch).items()} else: - feature_batch = features_val[i_sample:i_sample + batch_size] + feature_batch = features_val[i_sample : i_sample + batch_size] losses_val_batch = {k: l.numpy() for k, l in loss_eval_fn(feature_batch, batch).items()} for k, l in losses_val_batch.items(): losses_val[k] = losses_val.get(k, 0) + l * len(batch) diff --git a/models/utils.py b/models/utils.py index 98743b0..b8347d9 100644 --- a/models/utils.py +++ b/models/utils.py @@ -2,7 +2,7 @@ def epoch_from_name(name): - epoch, = re.findall(r'\d+', name) + (epoch,) = re.findall(r'\d+', name) return int(epoch) @@ -16,9 +16,7 @@ def latest_epoch(model_path): latest_gen_epoch = max(gen_epochs) latest_disc_epoch = max(disc_epochs) - assert ( - latest_gen_epoch == latest_disc_epoch - ), "Latest disc and gen epochs differ" + assert latest_gen_epoch == latest_disc_epoch, "Latest disc and gen epochs differ" return latest_gen_epoch diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..6156574 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,3 @@ +[tool.black] +line-length = 120 +skip-string-normalization = true \ No newline at end of file diff --git a/run_model_v4.py b/run_model_v4.py index 73ed5da..970fe06 100644 --- a/run_model_v4.py +++ b/run_model_v4.py @@ -45,9 +45,8 @@ def load_config(file): with open(file, 'r') as f: config = yaml.load(f, Loader=yaml.FullLoader) - assert ( - (config['feature_noise_power'] is None) == - (config['feature_noise_decay'] is None) + assert (config['feature_noise_power'] is None) == ( + config['feature_noise_decay'] is None ), 'Noise power and decay must be both provided' if 'lr_disc' not in config: @@ -113,44 +112,51 @@ def main(): for part in ['train', 'test']: evaluate_model( - model, path=prediction_path / part, - sample=( - (X_train, Y_train) if part == 'train' - else (X_test, Y_test) - ), - gen_sample_name=(None if part == 'train' else 'generated.dat') + model, + path=prediction_path / part, + sample=((X_train, Y_train) if part == 'train' else (X_test, Y_test)), + gen_sample_name=(None if part == 'train' else 'generated.dat'), ) else: features_noise = None if config['feature_noise_power'] is not None: + def features_noise(epoch): - current_power = config['feature_noise_power'] / (10**(epoch / config['feature_noise_decay'])) + current_power = config['feature_noise_power'] / (10 ** (epoch / config['feature_noise_decay'])) with writer_train.as_default(): tf.summary.scalar("features noise power", current_power, epoch) return current_power - save_model = SaveModelCallback( - model=model, path=model_path, save_period=config['save_every'] - ) + save_model = SaveModelCallback(model=model, path=model_path, save_period=config['save_every']) write_hist_summary = WriteHistSummaryCallback( - model, sample=(X_test, Y_test), - save_period=config['save_every'], writer=writer_val + model, sample=(X_test, Y_test), save_period=config['save_every'], writer=writer_val ) schedule_lr = ScheduleLRCallback( - model, writer=writer_val, + model, + writer=writer_val, func_gen=get_scheduler(config['lr_gen'], config['lr_schedule_rate_gen']), - func_disc=get_scheduler(config['lr_disc'], config['lr_schedule_rate_disc']) + func_disc=get_scheduler(config['lr_disc'], config['lr_schedule_rate_disc']), ) if continue_training: schedule_lr(next_epoch - 1) - train(Y_train, Y_test, model.training_step, model.calculate_losses, config['num_epochs'], - config['batch_size'], train_writer=writer_train, val_writer=writer_val, - callbacks=[schedule_lr, save_model, write_hist_summary], - features_train=X_train, features_val=X_test, features_noise=features_noise, - first_epoch=next_epoch) + train( + Y_train, + Y_test, + model.training_step, + model.calculate_losses, + config['num_epochs'], + config['batch_size'], + train_writer=writer_train, + val_writer=writer_val, + callbacks=[schedule_lr, save_model, write_hist_summary], + features_train=X_train, + features_val=X_test, + features_noise=features_noise, + first_epoch=next_epoch, + ) if __name__ == '__main__': From 72e9c5ead5d968c18f72d083ac46c5bae163404b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A1=D1=83=D1=85=D0=BE=D1=80=D0=BE=D1=81=D0=BE=D0=B2=20?= =?UTF-8?q?=D0=90=D0=BB=D0=B5=D0=BA=D1=81=D0=B5=D0=B9?= Date: Wed, 6 Oct 2021 23:05:28 +0300 Subject: [PATCH 5/6] specified black version --- .github/workflows/push.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index 9a65522..12ccb36 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -17,4 +17,6 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v2 - - uses: psf/black@stable \ No newline at end of file + - uses: psf/black@stable + with: + version: '19.10b0' \ No newline at end of file From d187191f5f07e5fda3e3d7d7b48f6d9aa00fdb18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A1=D1=83=D1=85=D0=BE=D1=80=D0=BE=D1=81=D0=BE=D0=B2=20?= =?UTF-8?q?=D0=90=D0=BB=D0=B5=D0=BA=D1=81=D0=B5=D0=B9?= Date: Thu, 14 Oct 2021 19:31:31 +0300 Subject: [PATCH 6/6] fixes --- .gitattributes | 3 +++ model_export/dump_graph.py | 3 ++- models/nn.py | 2 ++ 3 files changed, 7 insertions(+), 1 deletion(-) create mode 100644 .gitattributes diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..58f0190 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,3 @@ +* text=auto + +*.sh text eol=crlf diff --git a/model_export/dump_graph.py b/model_export/dump_graph.py index 4ff3ba3..02dd65b 100644 --- a/model_export/dump_graph.py +++ b/model_export/dump_graph.py @@ -7,7 +7,8 @@ from tensorflow.core.framework import attr_value_pb2 from tensorflow.python.tools import optimize_for_inference_lib -import tf2xla_pb2 + +from . import tf2xla_pb2 def model_to_graph( diff --git a/models/nn.py b/models/nn.py index 3f692ef..fb94b90 100644 --- a/models/nn.py +++ b/models/nn.py @@ -1,5 +1,7 @@ import tensorflow as tf +# DO NOT REMOVE. used by code inside eval() +import numpy as np # noqa: F401 custom_objects = {}