summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorJordan Gong <jordan.gong@protonmail.com>2021-02-19 21:03:42 +0800
committerJordan Gong <jordan.gong@protonmail.com>2021-02-19 21:03:42 +0800
commite33c22e556ed64e1c1fdb011d78a124d1489ad15 (patch)
tree4f11c8fd2a9653e71d92b210754a9db2eb4ef0be
parent4572598bca01d29b039f7d5291b08b53c1419528 (diff)
parenta31eb135c8cc3a8737fabca54fe3d5791f293753 (diff)
Merge branch 'data_parallel' into data_parallel_py3.8
-rw-r--r--models/model.py8
-rw-r--r--models/rgb_part_net.py2
2 files changed, 3 insertions, 7 deletions
diff --git a/models/model.py b/models/model.py
index ba87e3c..2ef3b80 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 538ca53..2af990e 100644
--- a/models/rgb_part_net.py
+++ b/models/rgb_part_net.py
@@ -86,7 +86,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)