run abl_all

cvae.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_abl_all_post import config
+from configs.iccv.config_abl_all import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
@@ -45,7 +45,7 @@ from models.cvae_model import CVAENew as CVAE
 
 
 def init_wandb():
-    run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
+    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae.py')
     artifact.add_file('models/cvae_model.py')
@@ -75,7 +75,7 @@ def kl_loss(mean, logvar):
 
 # 3d loss on posed !!!!!!!!!!!!
 
-# @tf.function
+@tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
@@ -478,7 +478,7 @@ def test():
 
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
@@ -528,7 +528,7 @@ if __name__ == '__main__':
     # test()
 
     ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20)
+    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/iccv/' + config['name'], max_to_keep=20)
 
 
     if config['pretrained_model'] != '':
@@ -543,9 +543,10 @@ if __name__ == '__main__':
 
         train_loss = train_step(batch)
 
-        data_valid = valid_iterator.get_next()
-        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
-        valid_loss = valid_step(data_valid)
+        if step % 10 == 0:
+            data_valid = valid_iterator.get_next()
+            data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
+            valid_loss = valid_step(data_valid)
 
         to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
         if step % 5000 == 0 or (step < 5000 and step % 500 == 0):