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

started training pyramid

Browse Source
This commit is contained in:
luni
2023-02-26 22:25:50 +00:00
parent a678724711
commit 64ecb900f5
8 changed files with 314 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cvae.py
+29 -13
View File
@@ -42,6 +42,7 @@ from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mes
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_model import CVAENew as CVAE
from custom_utils.custom_utils import Accumulator
def init_wandb():
@@ -529,7 +530,8 @@ if __name__ == '__main__':
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/iccv/' + config['name'], max_to_keep=20)
train_accumulator = Accumulator()
valid_accumulator = Accumulator()
if config['pretrained_model'] != '':
load_model(model, config['pretrained_model'])
@@ -542,31 +544,45 @@ 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)
if step % 10 == 0:
if step % 50 == 0:
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 }
# 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 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 balancing_activated_step == 0:
balancing_activated_step = step
g_balancing_coeff = g_balancing_coeff + 1 / 10000
to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
wandb.log(to_log, step = step)
if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
to_log = dict({ **train_loss, **valid_loss, **{ '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)
train_accumulator.reset()
valid_accumulator.reset()
print(step)