diff --git a/auto_encoder_body.py b/auto_encoder_body.py index dadd84f..b092832 100644 --- a/auto_encoder_body.py +++ b/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) - 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) train() \ No newline at end of file diff --git a/auto_encoder_curv.py b/auto_encoder_curv.py index b7d9703..04f2788 100644 --- a/auto_encoder_curv.py +++ b/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() diff --git a/auto_encoder_static.py b/auto_encoder_static.py index 0b22efc..3b02157 100644 --- a/auto_encoder_static.py +++ b/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() - # train() \ No newline at end of file + # test() + train() \ No newline at end of file diff --git a/configs/config_ae_body.py b/configs/config_ae_body.py index 89b2aa5..e2f9285 100644 --- a/configs/config_ae_body.py +++ b/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' diff --git a/configs/config_ae_curv.py b/configs/config_ae_curv.py index d00fe0c..461b1ba 100644 --- a/configs/config_ae_curv.py +++ b/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 diff --git a/configs/config_ae_static.py b/configs/config_ae_static.py index c59eb7b..00c5d82 100644 --- a/configs/config_ae_static.py +++ b/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), diff --git a/dataset.py b/dataset.py index e36f2eb..db20d92 100644 --- a/dataset.py +++ b/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])) diff --git a/models/base.py b/models/base.py index 3a149e2..fee7bc7 100644 --- a/models/base.py +++ b/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): diff --git a/models/cond_encoder.py b/models/cond_encoder.py index cd73c2b..10d1b5f 100644 --- a/models/cond_encoder.py +++ b/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):