update ae body

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]))
-    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'])
 
-    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,
+    e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones((pred.shape[0], 14475))))
-             'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2
+    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)
+
+    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, }
+  
 
 
 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 = valid_iterator.get_next()
-        # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
+        data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
-        valid_loss = {}#valid_step(data_valid)
+        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

@@ -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,\
+config['notes'] = 'convnext + convnext, add dense with tanh at the end of reverse convnext, | \
-    flattened bottleneck with 128 latent_dim\ remove global orientation\
+    flattened bottleneck with 128 latent_dim\ remove global orientation | \
-    fixed gender issue with deoriented bodies, again, aaand again\
+    fixed gender issue with deoriented bodies, again, aaand again | \
-    trained on 100 garments\
+    trained on 100 garments | \
-    remove multiplications with garment uv mask\
+    remove multiplications with garment uv mask | \
-    larger bottleneck, low res output upscaled\
+    larger bottleneck, low res output upscaled | \
-    larger batch size, 32 instead of 4\
+    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,6 +27,8 @@ 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]
@@ -33,3 +36,13 @@ config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae']['b
 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_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

@@ -2,11 +2,12 @@ config = dict()
 
 
 config['batch_size'] = 4
-config['train_dir'] = '/home/workspace/local-data/train_skirt_25_2'
+config['train_dir'] = '/home/workspace/local-data/train'
-config['valid_dir'] = '/home/workspace/local-data/train_skirt_25_2'
+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'] = config['cond_ae']
-config['cond_ae_body']['latent_dim'] = config['latent_dim']
+# config['cond_ae_body'] = dict()
-config['cond_ae_body']['kernel_size'] = 7
+# config['cond_ae_body']['latent_dim'] = config['latent_dim']
-config['cond_ae_body']['upscale_output'] = True
+# config['cond_ae_body']['kernel_size'] = 7
-config['cond_ae_body']['bottleneck_size'] = 512
+# config['cond_ae_body']['upscale_output'] = True
-config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
+# config['cond_ae_body']['bottleneck_size'] = 512
-config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
+# config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
-config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
+# config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
-config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
+# config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
-config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]
+# 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]