summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--models/model.py7
-rw-r--r--models/rgb_part_net.py2
2 files changed, 8 insertions, 1 deletions
diff --git a/models/model.py b/models/model.py
index 9748e46..0b8de50 100644
--- a/models/model.py
+++ b/models/model.py
@@ -150,6 +150,7 @@ class Model:
self.rgb_pn = RGBPartNet(self.in_channels, **model_hp,
image_log_on=self.image_log_on)
# Try to accelerate computation using CUDA or others
+ self.rgb_pn = nn.DataParallel(self.rgb_pn)
self.rgb_pn = self.rgb_pn.to(self.device)
self.optimizer = optim.Adam([
{'params': self.rgb_pn.ae.parameters(), **ae_optim_hp},
@@ -196,6 +197,12 @@ class Model:
# 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)
+ 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()
+ ))
loss = losses.sum()
loss.backward()
self.optimizer.step()
diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py
index 260eabd..845a477 100644
--- a/models/rgb_part_net.py
+++ b/models/rgb_part_net.py
@@ -85,7 +85,7 @@ class RGBPartNet(nn.Module):
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))
+ losses = (*losses, hpm_ba_trip, pn_ba_trip)
return losses, images
else:
return x.unsqueeze(1).view(-1)