From ab09d9bd0fe92d97f340feef1e2bbbcd33468953 Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Mon, 1 Mar 2021 19:10:15 +0800
Subject: Bug fixes

---
 models/model.py | 28 +++++++++++++++++-----------
 1 file changed, 17 insertions(+), 11 deletions(-)

(limited to 'models')

diff --git a/models/model.py b/models/model.py
index 2eeaf5e..70596ad 100644
--- a/models/model.py
+++ b/models/model.py
@@ -163,7 +163,7 @@ class Model:
                 )
             else:  # Different margins
                 self.triplet_loss = JointBatchTripletLoss(
-                    self.rgb_pn.hpm_num_parts,
+                    self.rgb_pn.module.hpm_num_parts,
                     triplet_is_hard, triplet_is_mean, triplet_margins
                 )
         else:  # Soft margins
@@ -226,13 +226,15 @@ class Model:
             embedding, ae_losses, images = self.rgb_pn(x_c1, x_c2)
             y = batch_c1['label'].to(self.device)
             # Duplicate labels for each part
-            y = y.repeat(self.rgb_pn.num_total_parts, 1)
-            trip_loss, dist, num_non_zero = self.triplet_loss(embedding, y)
+            y = y.repeat(self.rgb_pn.module.num_total_parts, 1)
+            trip_loss, dist, num_non_zero = self.triplet_loss(
+                embedding.contiguous(), y
+            )
             losses = torch.cat((
                 ae_losses.mean(0),
                 torch.stack((
-                    trip_loss[:self.rgb_pn.hpm_num_parts].mean(),
-                    trip_loss[self.rgb_pn.hpm_num_parts:].mean()
+                    trip_loss[:self.rgb_pn.module.hpm_num_parts].mean(),
+                    trip_loss[self.rgb_pn.module.hpm_num_parts:].mean()
                 ))
             ))
             loss = losses.sum()
@@ -254,28 +256,32 @@ class Model:
             # None-zero losses in batch
             if num_non_zero is not None:
                 self.writer.add_scalars('Loss/non-zero counts', {
-                    'HPM': num_non_zero[:self.rgb_pn.hpm_num_parts].mean(),
-                    'PartNet': num_non_zero[self.rgb_pn.hpm_num_parts:].mean()
+                    'HPM': num_non_zero[
+                           :self.rgb_pn.module.hpm_num_parts].mean(),
+                    'PartNet': num_non_zero[
+                               self.rgb_pn.module.hpm_num_parts:].mean()
                 }, self.curr_iter)
             # Embedding distance
-            mean_hpm_dist = dist[:self.rgb_pn.hpm_num_parts].mean(0)
+            mean_hpm_dist = dist[:self.rgb_pn.module.hpm_num_parts].mean(0)
             self._add_ranked_scalars(
                 'Embedding/HPM distance', mean_hpm_dist,
                 num_pos_pairs, num_pairs, self.curr_iter
             )
-            mean_pa_dist = dist[self.rgb_pn.hpm_num_parts:].mean(0)
+            mean_pa_dist = dist[self.rgb_pn.module.hpm_num_parts:].mean(0)
             self._add_ranked_scalars(
                 'Embedding/ParNet distance', mean_pa_dist,
                 num_pos_pairs, num_pairs, self.curr_iter
             )
             # Embedding norm
-            mean_hpm_embedding = embedding[:self.rgb_pn.hpm_num_parts].mean(0)
+            mean_hpm_embedding = embedding[
+                                 :self.rgb_pn.module.hpm_num_parts].mean(0)
             mean_hpm_norm = mean_hpm_embedding.norm(dim=-1)
             self._add_ranked_scalars(
                 'Embedding/HPM norm', mean_hpm_norm,
                 self.k, self.pr * self.k, self.curr_iter
             )
-            mean_pa_embedding = embedding[self.rgb_pn.hpm_num_parts:].mean(0)
+            mean_pa_embedding = embedding[
+                                self.rgb_pn.module.hpm_num_parts:].mean(0)
             mean_pa_norm = mean_pa_embedding.norm(dim=-1)
             self._add_ranked_scalars(
                 'Embedding/PartNet norm', mean_pa_norm,
-- 
cgit v1.2.3