more logging

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: