making ready for skirts

auto_encoder_curv.py

@@ -22,7 +22,7 @@ from models.cond_encoder import AutoEncoder
 testing = False
 def init_wandb():
     if not testing:
-        run = wandb.init(project='vae', name=config['name'], notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. same as 7')
+        run = wandb.init(project='vae', name=config['name'], notes=config['notes'])
         artifact = wandb.Artifact('code', type='code')
         artifact.add_file('auto_encoder_curv.py')
         artifact.add_file('losses.py')
@@ -79,7 +79,7 @@ def train_step(batch):
     with tf.GradientTape() as tape:
         pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)
 
-        ce_uv_loss = cce(batch['uv_curvature'], pred)
+        ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
         ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))
         loss = ce_uv_loss
 
@@ -166,7 +166,7 @@ if __name__ == '__main__':
     parser.add_argument('--name', type=str)
     args = vars(parser.parse_args())
 
-    cce = tf.keras.losses.CategoricalCrossentropy()
+    scce = tf.keras.losses.SparseCategoricalCrossentropy()
 
     resolution_modifier = ''
     if config['resolution'] != 512: