From a31eb135c8cc3a8737fabca54fe3d5791f293753 Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Fri, 19 Feb 2021 21:03:17 +0800 Subject: Correct cross reconstruction loss calculated in DataParallel --- models/model.py | 8 ++------ models/rgb_part_net.py | 2 +- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/models/model.py b/models/model.py index 4f24ac5..cac165e 100644 --- a/models/model.py +++ b/models/model.py @@ -198,12 +198,8 @@ class Model: # Duplicate labels for each part y = y.unsqueeze(1).repeat(1, self.rgb_pn.module.num_total_parts) losses, images = self.rgb_pn(x_c1, x_c2, y) - losses = torch.stack(( - # xrecon cano_cons pose_sim - losses[0].sum(), losses[1].mean(), losses[2].mean(), - # hpm_ba_trip pn_ba_trip - losses[3].mean(), losses[4].mean() - )) + # Combine losses from different data splits + losses = losses.mean() loss = losses.sum() loss.backward() self.optimizer.step() diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py index bd3a31c..67acac3 100644 --- a/models/rgb_part_net.py +++ b/models/rgb_part_net.py @@ -84,7 +84,7 @@ class RGBPartNet(nn.Module): pn_ba_trip = self.pn_ba_trip( x[self.hpm_num_parts:], y[self.hpm_num_parts:] ) - losses = (*losses, hpm_ba_trip, pn_ba_trip) + losses = torch.stack((*losses, hpm_ba_trip, pn_ba_trip)) return losses, images else: return x.unsqueeze(1).view(-1) -- cgit v1.2.3