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

save progress

Browse Source
This commit is contained in:
luni
2023-02-15 16:21:14 +00:00
parent d67865e7cd
commit a94be773d3
12 changed files with 656 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/cvae.py
+36 -40
View File
@@ -26,37 +26,25 @@ import numpy as np
import plotly.graph_objects as go
import plotly.express as ex
from PIL import Image
from postprocessing import fix_collision
from utils.postprocess import fix_collision
if config['is_skirt']:
from losses import edge_loss_skirt as edge_loss
from utils.losses import edge_loss_skirt as edge_loss
else:
from losses import edge_loss
from utils.losses import edge_loss
from losses import collision_ratio, kl_loss, normal_loss #, concept_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from utils.losses import collision_ratio, kl_loss, normal_loss, collision_loss
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
from dataset import DatasetFactory, StaticDataset
from data.dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_model import CVAENew as CVAE
from utils.utils import Accumulator
def init_wandb():
run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae.py')
artifact.add_file('models/cvae_model.py')
artifact.add_file('losses.py')
# artifact.add_file('configs/config_base.py')
# artifact.add_file('configs/config_cvae.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
artifact.add_file('dataset.py')
artifact.add_file('convnext/models/convnext_tf.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')]
config_file_names = [key for key in sys.modules.keys() if (key.startswith('configs.') or key.startswith('custom_utils.') or key.startswith('utils.') or key.startswith('data.') or key.startswith('models.') or key.startswith('scripts.'))]
for config_file_name in config_file_names:
file_name = '/'.join(config_file_name.split('.')) + '.py'
if os.path.exists(file_name):
@@ -101,7 +89,7 @@ def train_step(batch):
# is the output already masked?
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
@@ -159,7 +147,7 @@ def valid_step(batch):
# is the output already masked?
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
@@ -219,7 +207,7 @@ def test_step(batch):
# is the output already masked?
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1), axis=-1))
@@ -426,9 +414,6 @@ def sample(folder):
print()
def test():
str_split = config['name'].split('_')
test_idx = str_split.index('test')
@@ -518,6 +503,9 @@ if __name__ == '__main__':
in_memory=True,
is_cape=config['is_cape'])
train_accumulator = Accumulator()
valid_accumulator = Accumulator()
model = CVAE(config)
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
@@ -530,7 +518,6 @@ if __name__ == '__main__':
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20)
if config['pretrained_model'] != '':
load_model(model, config['pretrained_model'])
@@ -542,30 +529,39 @@ if __name__ == '__main__':
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
train_loss = train_step(batch)
train_accumulator.update(train_loss)
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)
valid_accumulator.update(valid_loss)
to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step % 5000 == 0:
manager.save(checkpoint_number=step)
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
if train_accumulator.get()['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
if balancing_activated_step == 0:
balancing_activated_step = step
g_balancing_coeff = g_balancing_coeff + 1 / 10000
g_balancing_coeff = g_balancing_coeff + 1 / 10000
to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
to_log = { **train_accumulator.get(), **valid_accumulator.get(), **{ 'current_step': step} }
if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
wandb.log(to_log, step = step)
wandb.log(to_log, step = step)
train_accumulator.reset()
valid_accumulator.reset()
if step % 5000 == 0:
manager.save(checkpoint_number=step)
print(step)