update ae curv
This commit is contained in:
+59
-33
@@ -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'])
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user