train non-pyramid with raw curv inputs

aaa_helpers.txt

@@ -1,5 +1,6 @@
 python merge_data.py -fi '/home/workspace/data/dataset/train' -fo '/home/workspace/local-data/train' -s 0 -e 1500
 nohup python auto_encoder_curv.py --name="lhunor_curv_raw_r2" > nohup/review/curv_raw_r2.out 2>&1 &
+nohup python cvae.py --name="lhunor_abl_all_raw_r2" > nohup/review/abl_all_raw_r2.out 2>&1 &
 
 
 one final convolutional layer for smoothing, on the final coordinates

configs/config_abl_all_post.py

@@ -19,9 +19,10 @@ config['only_curvatures'] =           False
 ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
 assert ablation_count == 1 or ablation_count == 0
 
-config['gpu'] = '1'
-config['name'] = 'al_all_post'
+config['gpu'] = '3'
+config['name'] = 'abl_all_raw_r2'
 config['notes'] = ''
+config['tags'] = ['review_2']
 
 config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
 if config['use_skinning']:
@@ -45,7 +46,7 @@ config['test_dir'] = '/home/workspace/data/dataset/test_ready'
 
 config['static_model'] = 'checkpoints/post/static_post_2'
 config['body_model'] = 'checkpoints/post/body_post_2'
-config['curvature_model'] = 'checkpoints/post/curv_post_1'
+config['curvature_model'] = 'checkpoints/post/curv_raw_r2'
 
 
 config['smpl_res'] = 14475

cvae.py

@@ -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='/wandb')
+    run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae.py')
     artifact.add_file('models/cvae_model.py')
@@ -58,7 +58,9 @@ def init_wandb():
     artifact.add_file('convnext/models/convnext_tf.py')
     config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
     for config_file_name in config_file_names:
-        artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
     run.log_artifact(artifact)
 
 uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False, is_smpl_orig=config['is_cape'])
@@ -73,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:
@@ -476,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())