update ae body
This commit is contained in:
+31
-23
@@ -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'])
|
||||
|
||||
Reference in New Issue
Block a user