From cf59a1f8339e937e31281907a61f484af99ce955 Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Mon, 15 Feb 2021 12:03:07 +0800 Subject: Add DataParallel support on new codebase --- models/model.py | 7 +++++++ 1 file changed, 7 insertions(+) (limited to 'models/model.py') diff --git a/models/model.py b/models/model.py index f79b832..532ded3 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() -- cgit v1.2.3