From feff6e315d2bb0d9c6d71c1c60770973a33aabd9 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 7 Oct 2022 21:55:36 +0000 Subject: [PATCH] update ae body --- auto_encoder_body.py | 54 ++++++++++++++++++++++----------------- configs/config_ae_body.py | 35 +++++++++++++++++-------- configs/config_base.py | 28 ++++++++++---------- 3 files changed, 70 insertions(+), 47 deletions(-) diff --git a/auto_encoder_body.py b/auto_encoder_body.py index 4c1bd81..cd2d9c3 100644 --- a/auto_encoder_body.py +++ b/auto_encoder_body.py @@ -2,13 +2,25 @@ import os from configs.config_ae_body import config os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] +import tensorflow as tf +gpus = tf.config.experimental.list_physical_devices('GPU') +if gpus: + try: + # Currently, memory growth needs to be the same across GPUs + for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + logical_gpus = tf.config.experimental.list_logical_devices('GPU') + print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs") + except RuntimeError as e: + # Memory growth must be set before GPUs have been initialized + print(e) + import time import argparse import wandb import random import cv2 import numpy as np -import tensorflow as tf from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch import plotly.graph_objects as go import plotly.express as ex @@ -21,7 +33,7 @@ from models.cond_encoder import AutoEncoder def init_wandb(): - run = wandb.init(project='vae', name=config['name'], notes=config['notes']) + run = wandb.init(project='ae', name=config['name'], notes=config['notes']) artifact = wandb.Artifact('code', type='code') artifact.add_file('auto_encoder_body.py') artifact.add_file('losses.py') @@ -50,33 +62,29 @@ def train_step(batch): loss = 1e3 * (1 - ms_ssim) - metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + metric = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((mesh_orig - mesh_pred) ** 2, axis=-1))) gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, } -# UPDATE! @tf.function def valid_step(batch): - pred = model([batch['uv_body_posed'], batch['uv_mask']]) + pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body], training=False) mesh_pred = uv2mesh(pred) mesh_orig = batch['mesh_body_posed'] - e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) - c_loss = 0 #1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])) - l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_body_posed'] - pred)), axis=1) * batch['uv_mask']) - ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed'], [0,2,3,1]) + 1, max_val=2)) - fft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_body_posed'], tf.zeros_like(batch['uv_body_posed'])))), axis=[1,2,3])) + + e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones((pred.shape[0], 14475)))) + l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred))), axis=1)) + ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + loss = 1e3 * (1 - ms_ssim) - mesh_orig_2 = uv2mesh(batch['uv_body_posed']) - metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) - metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) + + metric = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((mesh_orig - mesh_pred) ** 2, axis=-1))) + return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, } - return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, - 'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2 - } def plot_everything(batch: dict, type: str): @@ -115,13 +123,13 @@ def train(): batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) train_loss = train_step(batch) - # data_valid = valid_iterator.get_next() - # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic'])) - valid_loss = {}#valid_step(data_valid) + data_valid = valid_iterator.get_next() + data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic'])) + valid_loss = valid_step(data_valid) - if step % 2500 == 0: + if step % 5000 == 0 or (step < 5000 and step % 500 == 0): train_figs = plot_everything(batch, 'train') - valid_figs = {}#plot_everything(data_valid, 'valid') + valid_figs = plot_everything(data_valid, 'valid') wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step) else: wandb.log({ **train_loss, **valid_loss }, step=step) @@ -154,9 +162,9 @@ if __name__ == '__main__': factory = DatasetFactory() dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory']) - static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']]) + static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory']) - static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static']) + static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True) model = AutoEncoder(config['cond_ae']) optimizer = tf.optimizers.Adam(learning_rate=config['lr']) diff --git a/configs/config_ae_body.py b/configs/config_ae_body.py index 497cac0..449db77 100644 --- a/configs/config_ae_body.py +++ b/configs/config_ae_body.py @@ -1,19 +1,20 @@ from configs.config_base import config -config['gpu'] = '3' +config['gpu'] = '0' config['lr'] = 0.0001 -config['name'] = 'ae_body_10_3_skirt' +config['name'] = 'body_1' -config['notes'] = 'convnext + convnext, add dense with tanh at the end of reverse convnext,\ - flattened bottleneck with 128 latent_dim\ remove global orientation\ - fixed gender issue with deoriented bodies, again, aaand again\ - trained on 100 garments\ - remove multiplications with garment uv mask\ - larger bottleneck, low res output upscaled\ - larger batch size, 32 instead of 4\ +config['notes'] = 'convnext + convnext, add dense with tanh at the end of reverse convnext, | \ + flattened bottleneck with 128 latent_dim\ remove global orientation | \ + fixed gender issue with deoriented bodies, again, aaand again | \ + trained on 100 garments | \ + remove multiplications with garment uv mask | \ + larger bottleneck, low res output upscaled | \ + larger batch size, 32 instead of 4 | \ + calculate euclidean error on whole body, not just garment masked | \ ' -config['dataset_in_memory'] = True +config['dataset_in_memory'] = False config['batch_size'] = 32 @@ -26,10 +27,22 @@ config['signature_static'] = { 'mesh_mask': (14475,), } +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 config['cond_ae']['upscale_output'] = True config['cond_ae']['bottleneck_size'] = 512 config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3, 3] config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae']['bottleneck_size']] config['cond_ae']['decoder_input'] = (4, 2, config['cond_ae']['bottleneck_size']) config['cond_ae']['decoder_depths'] = [3, 3, 9, 3] -config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 128, 64, 16] \ No newline at end of file +config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 128, 64, 16] +# config['cond_ae_body'] = dict() +# config['cond_ae_body']['latent_dim'] = config['latent_dim'] +# config['cond_ae_body']['kernel_size'] = 7 +# config['cond_ae_body']['upscale_output'] = True +# config['cond_ae_body']['bottleneck_size'] = 512 +# config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3] +# config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']] +# config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size']) +# config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3] +# config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16] \ No newline at end of file diff --git a/configs/config_base.py b/configs/config_base.py index 90eadf7..a354bbf 100644 --- a/configs/config_base.py +++ b/configs/config_base.py @@ -2,11 +2,12 @@ config = dict() config['batch_size'] = 4 -config['train_dir'] = '/home/workspace/local-data/train_skirt_25_2' -config['valid_dir'] = '/home/workspace/local-data/train_skirt_25_2' +config['train_dir'] = '/home/workspace/local-data/train' +config['valid_dir'] = '/home/workspace/local-data/valid' config['use_filled'] = False config['resolution'] = 512 -config['is_skirt'] = True +config['is_skirt'] = False +config['dataset_in_memory'] = False config['latent_dim'] = 128 @@ -22,13 +23,14 @@ config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'] config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3] config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16] -config['cond_ae_body'] = dict() -config['cond_ae_body']['latent_dim'] = config['latent_dim'] -config['cond_ae_body']['kernel_size'] = 7 -config['cond_ae_body']['upscale_output'] = True -config['cond_ae_body']['bottleneck_size'] = 512 -config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3] -config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']] -config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size']) -config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3] -config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16] \ No newline at end of file +config['cond_ae_body'] = config['cond_ae'] +# config['cond_ae_body'] = dict() +# config['cond_ae_body']['latent_dim'] = config['latent_dim'] +# config['cond_ae_body']['kernel_size'] = 7 +# config['cond_ae_body']['upscale_output'] = True +# config['cond_ae_body']['bottleneck_size'] = 512 +# config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3] +# config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']] +# config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size']) +# config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3] +# config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16] \ No newline at end of file