more logging

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()

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()

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)
+    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)
+    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)
 
-    test()
+    # test()
-    # train()
+    train()

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

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:

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:

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:

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:

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: