diff options
| author | Jordan Gong <jordan.gong@protonmail.com> | 2021-01-12 20:37:48 +0800 |
|---|---|---|
| committer | Jordan Gong <jordan.gong@protonmail.com> | 2021-01-12 20:37:48 +0800 |
| commit | 37e307b960d565e7a58dfd22d0506c0f6f4fd122 (patch) | |
| tree | 5f7c3c35b8dbd793b3231596f7c8540d4375f09c /models | |
| parent | 36c65d17e29bcdbc27089d3b3eba177c8da549e8 (diff) | |
| parent | abb6989683829f12bf43ff19580444bf3396ac44 (diff) | |
Merge branch 'master' into python3.8
Diffstat (limited to 'models')
| -rw-r--r-- | models/model.py | 13 |
1 files changed, 7 insertions, 6 deletions
diff --git a/models/model.py b/models/model.py index d2e09fc..2c4e5a0 100644 --- a/models/model.py +++ b/models/model.py @@ -115,11 +115,11 @@ class Model: optim_hp = self.hp.get('optimizer', {}) sched_hp = self.hp.get('scheduler', {}) self.rgb_pn = RGBPartNet(self.train_size, self.in_channels, **model_hp) + # Try to accelerate computation using CUDA or others + self.rgb_pn = self._accelerate(self.rgb_pn) self.optimizer = optim.Adam(self.rgb_pn.parameters(), **optim_hp) self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, **sched_hp) self.writer = SummaryWriter(self._log_name) - # Try to accelerate computation using CUDA or others - self._accelerate() self.rgb_pn.train() # Init weights at first iter @@ -176,11 +176,12 @@ class Model: self.writer.close() break - def _accelerate(self): + def _accelerate(self, model: nn.Module) -> nn.Module: if not self.disable_acc: if torch.cuda.device_count() > 1: - self.rgb_pn = nn.DataParallel(self.rgb_pn) - self.rgb_pn = self.rgb_pn.to(self.device) + model = nn.DataParallel(model) + model = model.to(self.device) + return model def predict_all( self, @@ -204,7 +205,7 @@ class Model: model_hp = self.hp.get('model', {}) self.rgb_pn = RGBPartNet(ae_in_channels=self.in_channels, **model_hp) # Try to accelerate computation using CUDA or others - self._accelerate() + self.rgb_pn = self._accelerate(self.rgb_pn) self.rgb_pn.eval() gallery_samples, probe_samples = [], {} |