update

cvae.py

@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -35,10 +35,10 @@ from losses import collision_loss as collision_penalty
 from custom_utils.custom_utils import batch_to_dict, load_model
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
-from models.cvae import CVAE
+from models.cvae_model import CVAE
 
 
-name='cvae_15'
+name='cvae_16'
 def init_wandb():
     run = wandb.init(project='vae', name=name,
             notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \
@@ -47,9 +47,11 @@ offsets over average unposed garment. remove global orientation (in posed body i
 retrained body encoder with deoriented bodies \
 use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\
 fixed gender issue and retrained body encoder\
-remove shape from unposed garment')
+remove shape from unposed garment\
+training on 100 garments with offsets over the average of these')
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae.py')
+    artifact.add_file('models/cvae_model.py')
     artifact.add_file('dataset.py')
     artifact.add_file('losses.py')
     artifact.add_file('convnext/models/convnext_tf.py')
@@ -121,7 +123,7 @@ def kl_loss(mean, logvar):
     return kl_loss
 
 
-# @tf.function
+@tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
@@ -385,7 +387,7 @@ def sample(folder):
         print()
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
@@ -419,10 +421,10 @@ if __name__ == '__main__':
     static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
 
 
-    model = CVAE(avg_garment_path='custom_data/jumpsuit_20_avg_uv.npy',
-                 latent_dim=128, static_model='checkpoints/ae_static_3',
-                 body_model='checpoints/ae_body_6',
-                 curvature_model='checpoints/ae_curvature_6')
+    model = CVAE(avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
+                 latent_dim=128, static_model='checkpoints/ae_static_4',
+                 body_model='checpoints/ae_body_7',
+                 curvature_model='checpoints/ae_curvature_7')
 
 
     # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)