From 39393387993f62b247c987be70caba7522801112 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sun, 16 Oct 2022 14:54:33 +0000 Subject: [PATCH] more logging --- auto_encoder_body.py | 4 ++-- auto_encoder_curv.py | 4 ++-- auto_encoder_static.py | 17 ++++++++--------- cvae.py | 2 +- cvae_curv.py | 2 +- cvae_low.py | 2 +- cvae_pyramid.py | 8 +++++--- cvae_pyramid_old.py | 2 +- cvae_vae.py | 2 +- 9 files changed, 22 insertions(+), 21 deletions(-) diff --git a/auto_encoder_body.py b/auto_encoder_body.py index 4ae02f9..fbd2a13 100644 --- a/auto_encoder_body.py +++ b/auto_encoder_body.py @@ -173,8 +173,8 @@ if __name__ == '__main__': optimizer = tf.optimizers.Adam(learning_rate=config['lr']) ckpt_enc = tf.train.Checkpoint(model=model.encoder) - manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=3) + manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=20) ckpt_dec = tf.train.Checkpoint(model=model.decoder) - manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3) + manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=20) train() \ No newline at end of file diff --git a/auto_encoder_curv.py b/auto_encoder_curv.py index 1b015c2..15b629e 100644 --- a/auto_encoder_curv.py +++ b/auto_encoder_curv.py @@ -211,9 +211,9 @@ if __name__ == '__main__': optimizer = tf.optimizers.Adam(learning_rate=config['lr']) ckpt_enc = tf.train.Checkpoint(model=model.encoder) - manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=3) + manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=20) ckpt_dec = tf.train.Checkpoint(model=model.decoder) - manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3) + manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=20) if testing: test() diff --git a/auto_encoder_static.py b/auto_encoder_static.py index 9509d8f..a6044a7 100644 --- a/auto_encoder_static.py +++ b/auto_encoder_static.py @@ -56,7 +56,6 @@ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt']) @tf.function def train_step(batch): - with tf.GradientTape() as tape: pred = model([batch['uv_static'], batch['uv_mask']], training=True) @@ -131,7 +130,7 @@ def train(): data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic'])) valid_loss = valid_step(data_valid) - to_log = dict({ **train_loss, **valid_loss }) + to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } }) if step % 5000 == 0 or (step < 5000 and step % 500 == 0): train_figs = plot_everything(batch, 'train') valid_figs = plot_everything(data_valid, 'valid') @@ -186,7 +185,7 @@ def test(): if __name__ == '__main__': - # init_wandb() + init_wandb() allow_memory_growth() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) @@ -209,10 +208,10 @@ if __name__ == '__main__': model : AutoEncoder = AutoEncoder(config['cond_ae']) optimizer = tf.optimizers.Adam(learning_rate=config['lr']) - # ckpt_enc = tf.train.Checkpoint(model=model.encoder) - # manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=3) - # ckpt_dec = tf.train.Checkpoint(model=model.decoder) - # manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3) + ckpt_enc = tf.train.Checkpoint(model=model.encoder) + manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=20) + ckpt_dec = tf.train.Checkpoint(model=model.decoder) + manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=20) - test() - # train() \ No newline at end of file + # test() + train() \ No newline at end of file diff --git a/cvae.py b/cvae.py index 74527f0..2c1e058 100644 --- a/cvae.py +++ b/cvae.py @@ -388,7 +388,7 @@ if __name__ == '__main__': test() # ckpt = tf.train.Checkpoint(model=model) - # manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3) + # manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20) # valid_iterator = iter(dataset_valid) # step = 0 diff --git a/cvae_curv.py b/cvae_curv.py index 7a1b613..f42e0c7 100644 --- a/cvae_curv.py +++ b/cvae_curv.py @@ -388,7 +388,7 @@ if __name__ == '__main__': ckpt = tf.train.Checkpoint(model=model) - manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3) + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=20) step = 0 for batch in dataset_train: diff --git a/cvae_low.py b/cvae_low.py index ad8a9f9..fcb7bbb 100644 --- a/cvae_low.py +++ b/cvae_low.py @@ -391,7 +391,7 @@ if __name__ == '__main__': # sample('checkpoints/cvae_13') ckpt = tf.train.Checkpoint(model=model) - manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3) + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20) step = 0 for batch in dataset_train: diff --git a/cvae_pyramid.py b/cvae_pyramid.py index a7fc91c..024742b 100644 --- a/cvae_pyramid.py +++ b/cvae_pyramid.py @@ -21,7 +21,7 @@ import wandb import random import numpy as np import scipy as sp -from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces +from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces import plotly.graph_objects as go import plotly.express as ex from plotly.subplots import make_subplots @@ -498,7 +498,7 @@ if __name__ == '__main__': # sample('checkpoints/cvae_13') ckpt = tf.train.Checkpoint(model=model) - manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3) + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20) load_model(model, config['pyramid']['pretrained_pyramid']) model.static_enc.trainable = False @@ -538,7 +538,7 @@ if __name__ == '__main__': valid_loss = {} to_log = { **train_loss, **valid_loss, **{ 'current_step': step} } - if step % 2500 == 0: + if step % 5000 == 0 or (step < 5000 and step % 500 == 0): train_figs = plot_everything(batch, 'train') # valid_figs = plot_everything(data_valid, 'valid') valid_figs = {} @@ -546,6 +546,8 @@ if __name__ == '__main__': if step % 10000 == 0: manager.save(checkpoint_number=step) + if step == 0: + to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } if train_loss['train/euclid_error'] < 0.01: if balancing_activated_step == 0: diff --git a/cvae_pyramid_old.py b/cvae_pyramid_old.py index bbdbc10..5efc839 100644 --- a/cvae_pyramid_old.py +++ b/cvae_pyramid_old.py @@ -471,7 +471,7 @@ if __name__ == '__main__': ckpt = tf.train.Checkpoint(model=model) - manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3) + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=20) step = 0 for batch in dataset_train: diff --git a/cvae_vae.py b/cvae_vae.py index 44f7d7f..915ed68 100644 --- a/cvae_vae.py +++ b/cvae_vae.py @@ -490,7 +490,7 @@ if __name__ == '__main__': ckpt = tf.train.Checkpoint(model=model) - manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3) + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=20) step = 0 for batch in dataset_train: