From fc27a3cedf49a4751cef09f8d2bd56dedb1d8d57 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sun, 21 May 2023 09:59:57 +0000 Subject: [PATCH] add debug logging to cvae_train --- run.py | 14 ++++++++------ scripts/cvae.py | 43 ++++++++++++++++++++++++++++--------------- 2 files changed, 36 insertions(+), 21 deletions(-) diff --git a/run.py b/run.py index ca6fdbc..0e01c69 100644 --- a/run.py +++ b/run.py @@ -1,22 +1,24 @@ import os import sys -from configs.ae_curv_1 import config +from configs.abl_all import config import argparse -parser = argparse.ArgumentParser(description='') -parser.add_argument('--debug', action='store_true') -args = vars(parser.parse_args()) +# parser = argparse.ArgumentParser(description='') +# parser.add_argument('--debug', action='store_true') +# args = vars(parser.parse_args()) if __name__ == '__main__': - # import wandb - # wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838') + import wandb + # wandb.login(key='******************************************') + wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838') name = config['name'] script = config['script'] config = [key for key in sys.modules.keys() if key.startswith('configs.')][-1] current_file = '/'.join(__file__.split('/')[-2:]) + print(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/torch/{name}.out 2>&1 &') if args['debug']: os.system(f'python {script} --name lhunor_{name} --config {config}') else: diff --git a/scripts/cvae.py b/scripts/cvae.py index 8e9f118..fcd36f0 100644 --- a/scripts/cvae.py +++ b/scripts/cvae.py @@ -7,6 +7,7 @@ import argparse parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) parser.add_argument('--config', type=str) +parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) config_module = importlib.import_module(args['config']) @@ -128,8 +129,6 @@ def train_step(batch): } - - def valid_step(batch): if config['use_skinning']: mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed'] @@ -152,7 +151,6 @@ def valid_step(batch): else: mesh_pred_posed = mesh_pred - e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) @@ -221,6 +219,19 @@ def plot_everything(batch: dict, type: str): gt_posed = [batch['mesh_posed'].cpu().numpy()[i, mask[i], :] for i in range(mask.shape[0])] + # # plot garment_unposed, gt and mesh_posed next to eachother in plotly subplots + # from plotly.subplots import make_subplots + # fig = make_subplots(rows=1, cols=4, specs=[[{'type': 'surface'}, {'type': 'surface'}, {'type': 'surface'}, {'type': 'surface'}]]) + # fig.add_trace(go.Mesh3d(x=gt[0][:,0], y=gt[0][:,2], z=gt[0][:,1], i=faces[0][:,0], j=faces[0][:,1], k=faces[0][:,2], color='red', opacity=0.5), row=1, col=1) + # fig.add_trace(go.Mesh3d(x=garment_unposed[0][:,0], y=garment_unposed[0][:,2], z=garment_unposed[0][:,1], i=faces[0][:,0], j=faces[0][:,1], k=faces[0][:,2], color='blue', opacity=0.5), row=1, col=2) + # # add new subplot for gt_posed + # fig.add_trace(go.Mesh3d(x=gt_posed[0][:,0], y=gt_posed[0][:,2], z=gt_posed[0][:,1], i=faces[0][:,0], j=faces[0][:,1], k=faces[0][:,2], color='red', opacity=0.5), row=1, col=3) + # fig.add_trace(go.Mesh3d(x=mesh_posed[0][mask[0],0], y=mesh_posed[0][mask[0],2], z=mesh_posed[0][mask[0],1], i=faces[0][:,0], j=faces[0][:,1], k=faces[0][:,2], color='green', opacity=0.5), row=1, col=4) + # # make each subplot between -1 and 1 + # fig.update_layout(scene_aspectmode='cube', scene2_aspectmode='cube', scene3_aspectmode='cube', scene4_aspectmode='cube') + # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])), scene2=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])), scene3=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])), scene4=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1]))) + + # fig.show() mesh_posed_fixed = fix_collision(torch.tensor(mesh_posed, device=config['device']), batch['mesh_body_posed'], torch.tensor(mask, device=config['device'])).cpu().numpy() body = batch['mesh_body_posed'].cpu().numpy() @@ -412,9 +423,10 @@ def test(): def train(): + global g_balancing_coeff + global step valid_iterator = iter(valid_dataloader) - step = 0 for batch in train_dataloader: train_loss = train_step(batch) train_accumulator.update(train_loss) @@ -433,15 +445,14 @@ def train(): if step == 0: to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } - wandb.log(to_log, step=step) + if not args['debug']: + wandb.log(to_log, step=step) train_accumulator.reset() valid_accumulator.reset() print(step) 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 + g_balancing_coeff += 1 / 10000 if step % 10000 == 0: ckpt_saver(model, step, optimizer) @@ -450,11 +461,8 @@ def train(): if __name__ == '__main__': - init_wandb() - # parser = argparse.ArgumentParser(description='') - # parser.add_argument('--name', type=str) - # parser.add_argument('--debug', action='store_true') - # args = vars(parser.parse_args()) + if not args['debug']: + init_wandb() smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device']) skirt_modifier = '_skirt' if config['is_skirt'] else '' @@ -468,14 +476,19 @@ if __name__ == '__main__': model = CVAE(config) + optimizer = optim.Adam(model.parameters(), lr=config['lr']) _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3) - balancing_activated_step = 0 g_balancing_coeff = config['starting_balancing_coeff'] ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5) + step = 0 if config['pretrained_model'] != '': - load_model(model, config['pretrained_model']) + step = load_model(model, config['pretrained_model']) + if not args['debug']: + model = torch.compile(model) + wandb.watch(model, log='all', log_freq=1000, log_graph=True) + train() \ No newline at end of file