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

tried to profile code

Browse Source
This commit is contained in:
luni
2023-07-04 18:22:41 +02:00
parent 84b5af96cc
commit 9a81b82621
1 changed files with 30 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/cvae.py
+28 -6
View File
@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args()) args = vars(parser.parse_args())
if args['debug']: if args['debug']:
from configs.iccvw.abl_all import config from configs.iccvw.abl_all_4 import config
else: else:
config_module = importlib.import_module(args['config']) config_module = importlib.import_module(args['config'])
config = config_module.config config = config_module.config
@@ -39,6 +39,7 @@ else:
import torch import torch
import torch.nn as nn import torch.nn as nn
import torch.optim as optim import torch.optim as optim
from torch.profiler import profile, record_function, ProfilerActivity
from data.dataset import Dataset from data.dataset import Dataset
from models.cond_encoder import AutoEncoder from models.cond_encoder import AutoEncoder
from utils.checkpoints import CheckpointSaver, load_model from utils.checkpoints import CheckpointSaver, load_model
@@ -81,11 +82,13 @@ def kl_loss(mean: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
def train_step(batch): def train_step(batch):
# with profiler:
if config['use_skinning']: if config['use_skinning']:
mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed'] mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
else: else:
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed'] mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
# with record_function("forward"):
model.train() model.train()
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']]) pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
@@ -114,12 +117,14 @@ def train_step(batch):
if config['use_extra_losses']: if config['use_extra_losses']:
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff'] reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
loss = reconstr_loss + g_balancing_coeff * kl loss = g_balancing_coeff * kl + reconstr_loss
# with record_function("backward"):
optimizer.zero_grad() optimizer.zero_grad()
loss.backward() loss.backward()
optimizer.step() optimizer.step()
# with record_function("eval"):
model.eval() model.eval()
with torch.inference_mode(): with torch.inference_mode():
mean = torch.mean(mean) mean = torch.mean(mean)
@@ -127,7 +132,11 @@ def train_step(batch):
ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1)) ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric, # profiler.step()
return { 'train/loss': loss,
'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric,
'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio, 'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio,
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar, 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss, 'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
@@ -465,8 +474,8 @@ def train():
valid_accumulator.reset() valid_accumulator.reset()
print(step) print(step)
if train_loss['train/euclid_error'] < config['balancing_activation_limit']: # if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
g_balancing_coeff += 1 / 10000 # g_balancing_coeff += 1 / 10000
if step % 25000 == 0 and not args['debug']: if step % 25000 == 0 and not args['debug']:
ckpt_saver(model, step, optimizer, None) ckpt_saver(model, step, optimizer, None)
@@ -474,6 +483,7 @@ def train():
step += 1 step += 1
if __name__ == '__main__': if __name__ == '__main__':
if not args['debug']: if not args['debug']:
init_wandb() init_wandb()
@@ -499,10 +509,22 @@ if __name__ == '__main__':
step = 0 step = 0
if config['pretrained_model'] != '': if config['pretrained_model'] != '':
step = load_model(model, config['pretrained_model']) step = load_model(model, config['pretrained_model'], optimizer)
if not args['debug']: if not args['debug']:
# model = torch.compile(model) # model = torch.compile(model)
wandb.watch(model, log='all', log_freq=10000, log_graph=True) wandb.watch(model, log='all', log_freq=10000, log_graph=True)
# create a profiler object
# profiler = torch.profiler.profile(
# activities=[ProfilerActivity.CPU],
# schedule=torch.profiler.schedule(wait=2, warmup=1, active=3, repeat=2),
# on_trace_ready=torch.profiler.tensorboard_trace_handler("log_dir"),
# )
# with profiler:
# with record_function("model_inference"):
train() train()
# profiler.stop()
# print(profiler.key_averages().table(sort_by="cpu_time_total", row_limit=25))