From 929c48093c9f49a515420eb28d2678e48756b300 Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Sun, 14 Feb 2021 16:46:56 +0800
Subject: Prepare for DataParallel

---
 models/model.py        |  2 ++
 models/rgb_part_net.py | 12 +++++++++---
 utils/triplet_loss.py  |  2 --
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/models/model.py b/models/model.py
index 9cb6a8e..f79b832 100644
--- a/models/model.py
+++ b/models/model.py
@@ -193,6 +193,8 @@ class Model:
             x_c1 = batch_c1['clip'].to(self.device)
             x_c2 = batch_c2['clip'].to(self.device)
             y = batch_c1['label'].to(self.device)
+            # Duplicate labels for each part
+            y = y.unsqueeze(1).repeat(1, self.rgb_pn.num_total_parts)
             losses, images = self.rgb_pn(x_c1, x_c2, y)
             loss = losses.sum()
             loss.backward()
diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py
index 8ebcfd3..2aa680c 100644
--- a/models/rgb_part_net.py
+++ b/models/rgb_part_net.py
@@ -44,7 +44,8 @@ class RGBPartNet(nn.Module):
             ae_feature_channels * 2, out_channels, hpm_use_1x1conv,
             hpm_scales, hpm_use_avg_pool, hpm_use_max_pool
         )
-        empty_fc = torch.empty(self.hpm_num_parts + tfa_num_parts,
+        self.num_total_parts = self.hpm_num_parts + tfa_num_parts
+        empty_fc = torch.empty(self.num_total_parts,
                                out_channels, embedding_dims)
         self.fc_mat = nn.Parameter(empty_fc)
 
@@ -75,8 +76,13 @@ class RGBPartNet(nn.Module):
         x = self.fc(x)
 
         if self.training:
-            hpm_ba_trip = self.hpm_ba_trip(x[:self.hpm_num_parts], y)
-            pn_ba_trip = self.pn_ba_trip(x[self.hpm_num_parts:], y)
+            y = y.T
+            hpm_ba_trip = self.hpm_ba_trip(
+                x[:self.hpm_num_parts], y[:self.hpm_num_parts]
+            )
+            pn_ba_trip = self.pn_ba_trip(
+                x[self.hpm_num_parts:], y[self.hpm_num_parts:]
+            )
             losses = torch.stack((*losses, hpm_ba_trip, pn_ba_trip))
             return losses, images
         else:
diff --git a/utils/triplet_loss.py b/utils/triplet_loss.py
index d573ef4..954def2 100644
--- a/utils/triplet_loss.py
+++ b/utils/triplet_loss.py
@@ -9,9 +9,7 @@ class BatchAllTripletLoss(nn.Module):
         self.margin = margin
 
     def forward(self, x, y):
-        # Duplicate labels for each part
         p, n, c = x.size()
-        y = y.repeat(p, 1)
 
         # Euclidean distance p x n x n
         x_squared_sum = torch.sum(x ** 2, dim=2)
-- 
cgit v1.2.3