luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update ae body

Browse Source
This commit is contained in:
luni
2022-10-07 21:55:36 +00:00
parent 3c798b473d
commit feff6e315d
3 changed files with 70 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_body.py
+31 -23
View File
@@ -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'])
configs/config_ae_body.py
+24 -11
View File
@@ -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]
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]
configs/config_base.py
+15 -13
View File
@@ -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]
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]