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

add debug logging to cvae_train

Browse Source
This commit is contained in:
luni
2023-05-21 09:59:57 +00:00
parent 8255c1bbc9
commit fc27a3cedf
2 changed files with 36 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
run.py
+8 -6
View File
@@ -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:
scripts/cvae.py
+26 -13
View File
@@ -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)) } }
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__':
if not args['debug']:
init_wandb()
# parser = argparse.ArgumentParser(description='')
# parser.add_argument('--name', type=str)
# parser.add_argument('--debug', action='store_true')
# args = vars(parser.parse_args())
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()