From fdf36b27e7a80e83c0b4130d46a7877487ca8d99 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Mon, 26 Sep 2022 11:35:08 +0000 Subject: [PATCH] only load encoders of AEs --- configs/config_cvae.py | 21 +++++++++++++++++---- cvae.py | 19 ++++++------------- models/cvae_model.py | 36 +++++++++--------------------------- 3 files changed, 32 insertions(+), 44 deletions(-) diff --git a/configs/config_cvae.py b/configs/config_cvae.py index 56d84d1..1c5ecea 100644 --- a/configs/config_cvae.py +++ b/configs/config_cvae.py @@ -1,7 +1,20 @@ from configs.config_base import config -config['lr'] = 0.0001 +config['gpu'] = '3' config['name'] = 'cvae_19' +config['notes'] = '\ +offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \ +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, FINAL\ +remove shape from unposed garment\ +training on 100 garments with offsets over the average of these\ +fixed autoencoders!!!!!!!!!\ +cvae only loads encoders instead AEs\ +' + + +config['lr'] = 0.0001 config['starting_balancing_coeff'] = 0.0001 config['signature_dynamic'] = { 'id': (), @@ -19,9 +32,9 @@ config['signature_static'] = { 'uv_static': (512, 256, 3), } config['avg_garment_path'] = 'custom_data/train_100_no_shape_avg_uv.npy' -config['static_model'] = 'checkpoints/ae_static_6_copy' -config['body_model'] = 'checkpoints/ae_body_8_copy' -config['curvature_model'] = 'checkpoints/ae_curvature_7_copy' +config['static_model'] = 'checkpoints/ae_static_7_2_x' +config['body_model'] = 'checkpoints/ae_body_9_x' +config['curvature_model'] = 'checkpoints/ae_curv_9_x' config['freeze_encoders'] = True config['cvae'] = dict() diff --git a/cvae.py b/cvae.py index 30c82ba..205b10a 100644 --- a/cvae.py +++ b/cvae.py @@ -1,5 +1,6 @@ import os -os.environ["CUDA_VISIBLE_DEVICES"] = "0" +from configs.config_cvae import config +os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') @@ -35,24 +36,16 @@ from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d from dataset import DatasetFactory, StaticDataset from custom_utils.smpl import SMPL from models.cvae_model import CVAE -from configs.config_cvae import config def init_wandb(): run = wandb.init(project='vae', name=config['name'], - notes='\ -offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \ -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, FINAL\ -remove shape from unposed garment\ -training on 100 garments with offsets over the average of these\ -fixed autoencoders!!!!!!!!!') + notes=config['notes']) artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae.py') artifact.add_file('models/cvae_model.py') artifact.add_file('losses.py') - artifact.add_file('configs/config.py') + artifact.add_file('configs/config_base.py') artifact.add_file('configs/config_cvae.py') artifact.add_file('models/base.py') artifact.add_file('models/cond_encoder.py') @@ -76,7 +69,7 @@ def kl_loss(mean, logvar): return kl_loss -# @tf.function +@tf.function def train_step(batch): with tf.GradientTape() as tape: @@ -321,7 +314,7 @@ def sample(folder): if __name__ == '__main__': - # init_wandb() + init_wandb() allow_memory_growth() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) diff --git a/models/cvae_model.py b/models/cvae_model.py index c8e48b7..ddd13bc 100644 --- a/models/cvae_model.py +++ b/models/cvae_model.py @@ -8,19 +8,7 @@ import scipy as sp # import plotly.express as ex from custom_utils.custom_utils import load_model -from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block -from models.base import Bottleneck, VariationalEncoder, Decoder -from models.cond_encoder import AutoEncoder -# from models.ae_body_model import AutoEncoder as BodyAE -# from models.ae_curv_model import AutoEncoder as CurvatureAE -# from models.ae_static_model import AutoEncoder as StaticAE -# from models.ae_body_model import AutoEncoder32 as BodyAE32 -# from models.ae_curv_model import AutoEncoder32 as CurvatureAE32 -# from models.ae_static_model import AutoEncoder32 as StaticAE32 -# from models.ae_body_model import AutoEncoder128 as BodyAE128 -# from models.ae_curv_model import AutoEncoder128 as CurvatureAE128 -# from models.ae_static_model import AutoEncoder128 as StaticAE128 - +from models.base import AEEncoder, VariationalEncoder, Decoder class CVAE(tf.keras.Model): @@ -30,28 +18,23 @@ class CVAE(tf.keras.Model): self.latent_dim = config['latent_dim'] self.size = config['resolution'] - model : AutoEncoder = AutoEncoder(config['cond_ae']) + self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size']) if config['static_model'] != '': - load_model(model, config['static_model']) - # self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) - self.static_enc = model.encoder + load_model(self.static_enc, config['static_model'] + '_enc') - model : AutoEncoder = AutoEncoder(config['cond_ae']) + self.body_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size']) if config['body_model'] != '': - load_model(model, config['body_model']) - self.body_enc = model.encoder + load_model(self.body_enc, config['body_model'] + '_enc') - model : AutoEncoder = AutoEncoder(config['cond_ae']) + self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size']) if config['curvature_model'] != '': - load_model(model, config['curvature_model']) - self.curv_enc = model.encoder - + load_model(self.curv_enc, config['curvature_model'] + '_enc') + if config['freeze_encoders']: self.static_enc.trainable = False self.body_enc.trainable = False self.curv_enc.trainable = False - self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size']) self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size']) @@ -65,7 +48,6 @@ class CVAE(tf.keras.Model): tf.keras.layers.Dense(128, activation='relu'), tf.keras.layers.Dense(config['latent_dim']),]) - # self.shuffle = tf.keras.layers.Lambda(lambda x: tf.stack([tf.gather(x[i], tf.random.shuffle(tf.range(x[i].shape[0]))) for i in range(x.shape[0])])) # @tf.function def sample_single(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False): @@ -153,7 +135,7 @@ class CVAE(tf.keras.Model): return z_cond - def call(self, inputs, training=None, use_mask=False): + def call(self, inputs, training=None, use_mask=True): if use_mask: if self.size == 512: