tried to profile code

scripts/cvae.py

@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.abl_all import config
+    from configs.iccvw.abl_all_4 import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
@@ -39,6 +39,7 @@ else:
 import torch
 import torch.nn as nn
 import torch.optim as optim
+from torch.profiler import profile, record_function, ProfilerActivity
 from data.dataset import Dataset
 from models.cond_encoder import AutoEncoder
 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):
+    # with profiler:
     if config['use_skinning']:
         mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
     else:
         mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
 
+    # with record_function("forward"):
     model.train()
     pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
     mesh_pred = uv2mesh(pred)
@@ -114,20 +117,26 @@ def train_step(batch):
     if config['use_extra_losses']:
         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()
     loss.backward()
     optimizer.step()
 
+    # with record_function("eval"):
     model.eval()
     with torch.inference_mode():
         mean = torch.mean(mean)
         logvar = torch.mean(logvar)
         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'])
-    
-    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/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,
@@ -465,8 +474,8 @@ def train():
             valid_accumulator.reset()
             print(step)
 
-        if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
-            g_balancing_coeff += 1 / 10000 
+        # if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
+        #     g_balancing_coeff += 1 / 10000 
 
         if step % 25000 == 0 and not args['debug']:
             ckpt_saver(model, step, optimizer, None)
@@ -474,6 +483,7 @@ def train():
         step += 1
 
 
+
 if __name__ == '__main__':
     if not args['debug']:
         init_wandb()
@@ -499,10 +509,22 @@ if __name__ == '__main__':
 
     step = 0
     if config['pretrained_model'] != '':
-        step = load_model(model, config['pretrained_model'])
+        step = load_model(model, config['pretrained_model'], optimizer)
 
     if not args['debug']:
         # model = torch.compile(model)
         wandb.watch(model, log='all', log_freq=10000, log_graph=True)
 
-    train()
+    # 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()
+
+    # profiler.stop()
+    # print(profiler.key_averages().table(sort_by="cpu_time_total", row_limit=25))