From ece45943f851ef4d3f4a1040186870cddf3c57ef Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 7 Oct 2022 21:56:01 +0000 Subject: [PATCH] update ae curv --- auto_encoder_curv.py | 92 +++++++++++++++++++++++------------- configs/config_ae_curv.py | 9 ++-- custom_utils/custom_utils.py | 6 +-- 3 files changed, 66 insertions(+), 41 deletions(-) diff --git a/auto_encoder_curv.py b/auto_encoder_curv.py index fe752d4..393fc64 100644 --- a/auto_encoder_curv.py +++ b/auto_encoder_curv.py @@ -1,15 +1,28 @@ import os from configs.config_ae_curv 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 numpy as np -import tensorflow as tf -from custom_utils.custom_utils import UV2Mesh, maskFacesBatch +from custom_utils.custom_utils import UV2Mesh, UV2MeshLabel, maskFacesBatch import plotly.graph_objects as go import plotly.express as ex @@ -22,7 +35,7 @@ from models.cond_encoder import AutoEncoder testing = False def init_wandb(): if not testing: - 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_curv.py') artifact.add_file('losses.py') @@ -34,7 +47,7 @@ def init_wandb(): artifact.add_file('convnext/models/convnext_tf.py') run.log_artifact(artifact) -uv2mesh = UV2Mesh(res=256) +uv2mesh = UV2MeshLabel(res=256, is_skirt=config['is_skirt']) g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7) @@ -43,6 +56,12 @@ def get_mean_iou(y_true, y_pred): g_mean_iou.update_state(y_true, y_pred) return g_mean_iou.result() +g_mean_iou_mesh = tf.keras.metrics.MeanIoU(num_classes=7) +def get_mean_iou_mesh(y_true, y_pred): + g_mean_iou_mesh.reset_state() + g_mean_iou_mesh.update_state(y_true, y_pred) + return g_mean_iou_mesh.result() + # https://github.com/NifTK/NiftyNet/blob/dev/niftynet/layer/loss_segmentation.py def dice_dense(prediction, ground_truth, weight_map=None): """ @@ -75,7 +94,6 @@ def dice_dense(prediction, ground_truth, weight_map=None): @tf.function def train_step(batch): - with tf.GradientTape() as tape: pred = model([batch['uv_curvature'], batch['uv_mask']], training=True) @@ -83,54 +101,62 @@ def train_step(batch): ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1)) loss = ce_uv_loss + pred_mesh = tf.math.argmax(uv2mesh(pred), axis=-1) mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1)) + mean_iou_mesh = get_mean_iou(tf.cast(batch['mesh_curvature'], tf.float32) * batch['mesh_mask'], tf.cast(pred_mesh, tf.float32) * batch['mesh_mask']) dice = dice_dense(pred, batch['uv_curvature']) - + gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) - return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou, 'train/dice': dice } + return { 'train/loss': loss, 'train/ce_uv_loss': ce_uv_loss, 'train/ssim': ms_ssim, 'train/mean_iou': mean_iou, 'train/mean_iou_mesh': mean_iou_mesh, 'train/dice': dice } @tf.function def valid_step(batch): - uv_curvature = tf.one_hot(batch['uv_curvature'], 7) + pred = model([batch['uv_curvature'], batch['uv_mask']], training=False) - pred = model([uv_curvature, batch['uv_mask']]) + ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred) + ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1)) + loss = ce_uv_loss - ce_uv_loss = cce(batch['uv_curvature'], pred) - ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1)) - loss = ce_uv_loss + pred_mesh = tf.math.argmax(uv2mesh(pred), axis=-1) + mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1)) + mean_iou_mesh = get_mean_iou(tf.cast(batch['mesh_curvature'], tf.float32) * batch['mesh_mask'], tf.cast(pred_mesh, tf.float32) * batch['mesh_mask']) + dice = dice_dense(pred, batch['uv_curvature']) - mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1)) - dice = dice_dense(pred, uv_curvature) + return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou, 'valid/mean_iou_mesh': mean_iou_mesh, 'valid/dice': dice } + - return { 'valid/loss': loss, 'valid/ce_uv_loss': ce_uv_loss, 'valid/ssim': ms_ssim, 'valid/mean_iou': mean_iou, 'valid/dice': dice #'ce_mesh_loss': ce_mesh_loss - } +def plot_everything(batch: dict, type: str): + pred = model([batch['uv_curvature'], batch['uv_mask']], training=False) + + orig_img = tf.math.argmax(batch['uv_curvature'], axis=-1, output_type=tf.dtypes.int32) + pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32) + img = tf.concat((orig_img, pred_img), axis=2).numpy() + imgs = { type + '/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(img.shape[0])] } + + return imgs def train(): + valid_iterator = iter(dataset_valid) + step = 0 for batch in dataset_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, signature_dynamic)) - valid_loss = {}#valid_step(data_valid) - data_valid = 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) to_log = { **train_loss, **valid_loss} - if step % 2500 == 0: - pred = model([data_valid['uv_curvature'], data_valid['uv_mask']], training=False) - - orig_img = tf.math.argmax(data_valid['uv_curvature'], axis=-1, output_type=tf.dtypes.int32) - pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32) - img = tf.concat((orig_img, pred_img), axis=2).numpy() - # imgs = { 'valid/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(4)] } - imgs = { 'train/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(img.shape[0])] } - to_log = { **to_log, **imgs} + if step % 5000 == 0 or (step < 5000 and step % 500 == 0): + imgs_train = plot_everything(batch, 'train') + imgs_valid = plot_everything(data_valid, 'valid') + to_log = { **to_log, **imgs_train, **imgs_valid } wandb.log(to_log, step=step) print(step) @@ -174,10 +200,10 @@ 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']]) - static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']]) - dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic']) - static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static']) + 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']], 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'], 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_curv.py b/configs/config_ae_curv.py index a58b339..051c2ae 100644 --- a/configs/config_ae_curv.py +++ b/configs/config_ae_curv.py @@ -1,13 +1,12 @@ from configs.config_base import config -config['gpu'] = '0' +config['gpu'] = '2' config['lr'] = 0.001 -config['name'] = 'ae_curv_9_skirt' #'ae_curv_9' -config['notes'] = 'convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. switched back to sparse CCE, normal one was giving NaN' +config['name'] = 'curv_1_3' #'ae_curv_9' +config['notes'] = 'convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. switched back to sparse CCE, normal one was giving NaN | \ + try ssim loss' config['cond_ae']['is_label'] = True -config['cond_ae']['upscale_output'] = False - config['signature_dynamic'] = { 'id': (), 'uv_curvature': (512, 256, 7), diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index a94c9b5..d5a51b4 100644 --- a/custom_utils/custom_utils.py +++ b/custom_utils/custom_utils.py @@ -71,12 +71,12 @@ class UV2Mesh(): class UV2MeshLabel(): - def __init__(self, res=256): + def __init__(self, res=256, is_skirt=False): super(UV2MeshLabel, self).__init__() self.res = res - uv_map = np.load('custom_utils/stretched_uv_map_3.npy') + uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + '.npy') self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32) - m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item() + m2uv = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + '.npy', allow_pickle=True).item() m2uv = m2uv_array(m2uv) self.uv_pixels = self.uv_map_discreet[m2uv]