working auto encoders

auto_encoder_curv.py

@@ -1,11 +1,12 @@
 import os
+from configs.config_ae_curv import config
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import time
 import argparse
 import wandb
 import random
 import numpy as np
 import tensorflow as tf
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
 
 from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
@@ -16,7 +17,6 @@ from tf_utils import check_tf_version, allow_memory_growth
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict
 from models.cond_encoder import AutoEncoder
-from configs.config_ae_curv import config
 
 
 testing = False
@@ -26,7 +26,7 @@ def init_wandb():
         artifact = wandb.Artifact('code', type='code')
         artifact.add_file('auto_encoder_curv.py')
         artifact.add_file('losses.py')
-        artifact.add_file('configs/config.py')
+        artifact.add_file('configs/config_base.py')
         artifact.add_file('configs/config_ae_curv.py')
         artifact.add_file('models/base.py')
         artifact.add_file('models/cond_encoder.py')
@@ -37,45 +37,6 @@ def init_wandb():
 uv2mesh = UV2Mesh(res=256)
 
 
-def plot_3d(mesh_gt, mesh_pred, faces):
-    mesh_pred[:,1] += 2
-    fig = go.Figure(data=[
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_gt[:,0],
-            y=mesh_gt[:,1],
-            z=mesh_gt[:,2],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2]
-        ),
-        go.Mesh3d(
-            # 8 vertices of a cube
-            x=mesh_pred[:,0],
-            y=mesh_pred[:,1],
-            z=mesh_pred[:,2],
-            # i, j and k give the vertices of triangles
-            i = faces[:,0],
-            j = faces[:,1],
-            k = faces[:,2],
-        )
-    ])
-
-    fig.update_layout(scene_aspectmode='manual',
-                    scene_aspectratio=dict(x=1, y=2, z=1))
-    fig.update_layout(
-        scene = dict(
-            xaxis = dict(nticks=4, range=[-1,1],),
-                        yaxis = dict(nticks=4, range=[-1,3],),
-                        zaxis = dict(nticks=4, range=[-1,1],),)
-                        )
-
-    return fig
-
-
-
-
 g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7)
 def get_mean_iou(y_true, y_pred):
     g_mean_iou.reset_state()
@@ -112,13 +73,13 @@ def dice_dense(prediction, ground_truth, weight_map=None):
     return tf.reduce_mean(dice_score)
 
 
-# @tf.function
+@tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
         uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
 
-        pred = model([uv_curvature, batch['uv_mask']])
+        pred = model([uv_curvature, batch['uv_mask']], training=True)
 
 
         ce_uv_loss = scce(batch['uv_curvature'], pred)
@@ -166,20 +127,21 @@ def train():
         to_log = { **train_loss, **valid_loss}
         if step % 2500 == 0:
             uv_curvature = tf.one_hot(data_valid['uv_curvature'], 7)
-            pred = model([uv_curvature, data_valid['uv_mask']])
+            pred = model([uv_curvature, data_valid['uv_mask']], training=False)
             
             orig_img = data_valid['uv_curvature']
             pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)
             img = tf.concat((orig_img, pred_img), axis=2).numpy()
             # imgs = { 'valid/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(4)] }
-            imgs = { 'train/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(4)] }
+            imgs = { 'train/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(img.shape[0])] }
             to_log = { **to_log, **imgs}
 
         wandb.log(to_log, step=step)  
         print(step)
 
-        if step % 2500 == 0:
-            manager.save(checkpoint_number=step)
+        if step % 10000 == 0:
+            manager_enc.save(checkpoint_number=step)
+            manager_dec.save(checkpoint_number=step)
 
         step += 1
 
@@ -201,7 +163,7 @@ def test():
 
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
@@ -223,9 +185,10 @@ if __name__ == '__main__':
     model = AutoEncoder(config['cond_ae'])
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
 
-    ckpt = tf.train.Checkpoint(model=model)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], 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=3)
+    ckpt_dec = tf.train.Checkpoint(model=model.decoder)
+    manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3)
 
     if testing:
         test()