working auto encoders

auto_encoder_body.py

@@ -31,7 +31,7 @@ def init_wandb():
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('auto_encoder_body.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_body.py')
     artifact.add_file('models/base.py')
     artifact.add_file('models/cond_encoder.py')
@@ -45,7 +45,7 @@ uv2mesh = UV2Mesh(res=256)
 @tf.function
 def train_step(batch):
     with tf.GradientTape() as tape:
-        pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body])
+        pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body], training=True)
 
         mesh_pred = uv2mesh(pred)
         mesh_orig = batch['mesh_body_posed']
@@ -87,7 +87,7 @@ def valid_step(batch):
 
 def plot_everything(batch: dict, type: str):
 
-    pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body])
+    pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body], training=False)
     
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_body_posed']
@@ -133,7 +133,8 @@ def train():
             wandb.log({ **train_loss, **valid_loss }, step=step)
             
         if step % 10000 == 0:
-            manager.save(checkpoint_number=step)
+            manager_enc.save(checkpoint_number=step)
+            manager_dec.save(checkpoint_number=step)
 
         print(step)
 
@@ -141,7 +142,7 @@ def train():
 
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
@@ -165,7 +166,9 @@ if __name__ == '__main__':
     model = AutoEncoder(config['cond_ae'])
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
 
-    ckpt = tf.train.Checkpoint(model=model)
+    ckpt_enc = tf.train.Checkpoint(model=model.encoder)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
+    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)
 
     train()

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:
+        if step % 10000 == 0:
-            manager.save(checkpoint_number=step)
+            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)
+    ckpt_enc = tf.train.Checkpoint(model=model.encoder)
-    manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
+    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()

auto_encoder_static.py

@@ -83,7 +83,7 @@ def valid_step(batch):
 
 
 def plot_everything(batch: dict, type: str):
-    pred = model([batch['uv_static'], batch['uv_mask']])
+    pred = model([batch['uv_static'], batch['uv_mask']], training=False)
     
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_static']
@@ -151,7 +151,7 @@ def test():
 
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)
@@ -178,5 +178,5 @@ if __name__ == '__main__':
     ckpt_dec = tf.train.Checkpoint(model=model.decoder)
     manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3)
 
-    test()
+    # test()
-    # train()
+    train()

configs/config_ae_body.py

@@ -1,6 +1,6 @@
 from configs.config_base import config
 
-config['gpu'] = '1'
+config['gpu'] = '0'
 config['lr'] = 0.0001
 config['name'] = 'ae_body_9'
 

configs/config_ae_curv.py

@@ -1,7 +1,8 @@
 from configs.config_base import config
 
+config['gpu'] = '2'
 config['lr'] = 0.001
-config['name'] = 'ae_body_9'
+config['name'] = 'ae_curv_9'
 
 config['cond_ae']['is_label'] = True
 

configs/config_ae_static.py

@@ -1,8 +1,8 @@
 from configs.config_base import config
 
-config['gpu'] = '0'
+config['gpu'] = '1'
 config['lr'] = 0.0005
-config['name'] = 'ae_static_7_x'
+config['name'] = 'ae_static_7_2'
 
 config['signature_dynamic'] = { 'id': (), }
 config['signature_static'] = { 'uv_static': (512, 256, 3),

dataset.py

@@ -54,7 +54,7 @@ from custom_utils.custom_utils import nest_conditions, batch_to_dict
 
 
 class StaticDataset():
-    def __init__(self, folder, signature, use_filled=True, resolutions=[512]) -> None:
+    def __init__(self, folder, signature, use_filled=False, resolutions=[512]) -> None:
         self.store = dict()
         self.signature = signature
         self.folder = folder
@@ -139,7 +139,7 @@ class DatasetFactory():
         return tf_signature
 
     
-    def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=True, resolutions=[512]):
+    def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512]):
         # assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
         assert set(resolutions).issubset(set([32, 128, 512]))
 

models/base.py

@@ -64,7 +64,10 @@ class AEEncoder(tf.keras.Model):
         x = self.encoder(x, training=training)
         x = self.bottleneck(x, training=training)
 
-        return x
+        if len(inputs) == 2:
+            return x, inputs[1]
+        else:
+            return x
 
 
 class Decoder(tf.keras.Model):

models/cond_encoder.py

@@ -8,7 +8,7 @@ class AutoEncoder(tf.keras.Model):
         super(AutoEncoder, self).__init__()
 
         self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
-        self.decoder : Decoder= Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
+        self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
 
 
     def call(self, inputs, training=None):