diff options
| -rw-r--r-- | models/model.py | 13 |
1 files changed, 8 insertions, 5 deletions
diff --git a/models/model.py b/models/model.py index 90d48e0..a086e7b 100644 --- a/models/model.py +++ b/models/model.py @@ -157,13 +157,15 @@ class Model: self.rgb_pn.hpm_num_parts, triplet_margins ) # 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.ba_triplet_loss = nn.DataParallel(self.ba_triplet_loss) self.ba_triplet_loss = self.ba_triplet_loss.to(self.device) self.optimizer = optim.Adam([ - {'params': self.rgb_pn.ae.parameters(), **ae_optim_hp}, - {'params': self.rgb_pn.pn.parameters(), **pn_optim_hp}, - {'params': self.rgb_pn.hpm.parameters(), **hpm_optim_hp}, - {'params': self.rgb_pn.fc_mat, **fc_optim_hp} + {'params': self.rgb_pn.module.ae.parameters(), **ae_optim_hp}, + {'params': self.rgb_pn.module.pn.parameters(), **pn_optim_hp}, + {'params': self.rgb_pn.module.hpm.parameters(), **hpm_optim_hp}, + {'params': self.rgb_pn.module.fc_mat, **fc_optim_hp} ], **optim_hp) sched_gamma = sched_hp.get('gamma', 0.9) sched_step_size = sched_hp.get('step_size', 500) @@ -206,7 +208,7 @@ class Model: y = y.repeat(self.rgb_pn.num_total_parts, 1) triplet_loss = self.ba_triplet_loss(feature, y) losses = torch.cat(( - ae_losses, + ae_losses.mean(0), torch.stack(( triplet_loss[:self.rgb_pn.hpm_num_parts].mean(), triplet_loss[self.rgb_pn.hpm_num_parts:].mean() @@ -318,6 +320,7 @@ class Model: model_hp.pop('triplet_margins', None) self.rgb_pn = RGBPartNet(self.in_channels, self.in_size, **model_hp) # 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.rgb_pn.eval() |