more logging

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