From 31be2688653a94c3d8f4e707a118bd311803cc73 Mon Sep 17 00:00:00 2001
From: centos Cloud User <centos@dnadna-elliot.novalocal>
Date: Wed, 24 Jan 2024 10:14:02 +0100
Subject: [PATCH] debugging cosypose training

---
 .../cosypose/cosypose/training/train_pose.py  | 70 ++++++++++---------
 happypose/toolbox/lib3d/transform.py          |  5 +-
 2 files changed, 42 insertions(+), 33 deletions(-)

diff --git a/happypose/pose_estimators/cosypose/cosypose/training/train_pose.py b/happypose/pose_estimators/cosypose/cosypose/training/train_pose.py
index 39163c1a..01040496 100644
--- a/happypose/pose_estimators/cosypose/cosypose/training/train_pose.py
+++ b/happypose/pose_estimators/cosypose/cosypose/training/train_pose.py
@@ -422,44 +422,50 @@ def train_epoch():
             iterator = tqdm(ds_iter_train, ncols=80)
             t = time.time()
             for n, sample in enumerate(iterator):
-                if n > 0:
-                    meters_time["data"].add(time.time() - t)
-
-                optimizer.zero_grad()
-
-                t = time.time()
-                loss = h(data=sample, meters=meters_train)
-                meters_time["forward"].add(time.time() - t)
-                iterator.set_postfix(loss=loss.item())
-                meters_train["loss_total"].add(loss.item())
-
-                t = time.time()
-                loss.backward()
-                total_grad_norm = torch.nn.utils.clip_grad_norm_(
-                    model.parameters(),
-                    max_norm=args.clip_grad_norm,
-                    norm_type=2,
-                )
-                meters_train["grad_norm"].add(torch.as_tensor(total_grad_norm).item())
-
-                optimizer.step()
-                meters_time["backward"].add(time.time() - t)
-                meters_time["memory"].add(
-                    torch.cuda.max_memory_allocated() / 1024.0**2,
-                )
-
-                if epoch < args.n_epochs_warmup:
-                    lr_scheduler_warmup.step()
-                t = time.time()
+                if n < 5:
+                    if n > 0:
+                        meters_time["data"].add(time.time() - t)
+
+                    optimizer.zero_grad()
+
+                    t = time.time()
+                    loss = h(data=sample, meters=meters_train)
+                    meters_time["forward"].add(time.time() - t)
+                    iterator.set_postfix(loss=loss.item())
+                    meters_train["loss_total"].add(loss.item())
+
+                    t = time.time()
+                    loss.backward()
+                    total_grad_norm = torch.nn.utils.clip_grad_norm_(
+                        model.parameters(),
+                        max_norm=args.clip_grad_norm,
+                        norm_type=2,
+                    )
+                    meters_train["grad_norm"].add(torch.as_tensor(total_grad_norm).item())
+
+                    optimizer.step()
+                    meters_time["backward"].add(time.time() - t)
+                    meters_time["memory"].add(
+                        torch.cuda.max_memory_allocated() / 1024.0**2,
+                    )
+
+                    if epoch < args.n_epochs_warmup:
+                        lr_scheduler_warmup.step()
+                    t = time.time()
+                else:
+                    continue
             if epoch >= args.n_epochs_warmup:
                 lr_scheduler.step()
 
         @torch.no_grad()
         def validation():
             model.eval()
-            for sample in tqdm(ds_iter_val, ncols=80):
-                loss = h(data=sample, meters=meters_val)
-                meters_val["loss_total"].add(loss.item())
+            for n, sample in enumerate(tqdm(ds_iter_val, ncols=80)):
+                if n < 5:
+                    loss = h(data=sample, meters=meters_val)
+                    meters_val["loss_total"].add(loss.item())
+                else:
+                    continue
 
         @torch.no_grad()
         def test():
diff --git a/happypose/toolbox/lib3d/transform.py b/happypose/toolbox/lib3d/transform.py
index 001a19c2..281f3ae2 100644
--- a/happypose/toolbox/lib3d/transform.py
+++ b/happypose/toolbox/lib3d/transform.py
@@ -74,7 +74,10 @@ def __init__(
                 else:
                     rotation_np = rotation
             else:
-                raise ValueError
+                if isinstance(rotation, list):
+                    rotation_np = np.array(rotation)
+                else:
+                    raise ValueError
 
             if rotation_np.size == 4:
                 quaternion_xyzw = rotation_np.flatten().tolist()