only load encoders of AEs
This commit is contained in:
+17
-4
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
+8
-26
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user