diff --git a/ppgan/engine/trainer.py b/ppgan/engine/trainer.py index 30d2ccda1efead5b486628fc65e7874b33473d1e..756cbfc9c0288fcfc412c216b88aa1f7f01e762f 100755 --- a/ppgan/engine/trainer.py +++ b/ppgan/engine/trainer.py @@ -79,6 +79,7 @@ class Trainer: self.max_eval_steps = cfg.model.get('max_eval_steps', None) self.local_rank = ParallelEnv().local_rank + self.world_size = ParallelEnv().nranks self.log_interval = cfg.log_config.interval self.visual_interval = cfg.log_config.visiual_interval self.weight_interval = cfg.snapshot_config.interval @@ -217,8 +218,7 @@ class Trainer: def test(self): if not hasattr(self, 'test_dataloader'): self.test_dataloader = build_dataloader(self.cfg.dataset.test, - is_train=False, - distributed=False) + is_train=False) iter_loader = IterLoader(self.test_dataloader) if self.max_eval_steps is None: self.max_eval_steps = len(self.test_dataloader) @@ -231,6 +231,10 @@ class Trainer: self.model.setup_train_mode(is_train=False) for i in range(self.max_eval_steps): + if self.max_eval_steps < self.log_interval or i % self.log_interval == 0: + self.logger.info('Test iter: [%d/%d]' % + (i * self.world_size, self.max_eval_steps * self.world_size)) + data = next(iter_loader) self.model.setup_input(data) self.model.test_iter(metrics=self.metrics) @@ -264,9 +268,6 @@ class Trainer: step=self.batch_id, is_save_image=True) - if i % self.log_interval == 0: - self.logger.info('Test iter: [%d/%d]' % - (i, self.max_eval_steps)) if self.metrics: for metric_name, metric in self.metrics.items(): diff --git a/ppgan/metrics/lpips.py b/ppgan/metrics/lpips.py index aadeabfa53fdfd4ed9d68b3d415747f4935b517f..2a77eaf1b59de247bacf91b31cd79278574411df 100644 --- a/ppgan/metrics/lpips.py +++ b/ppgan/metrics/lpips.py @@ -89,6 +89,12 @@ class LPIPSMetric(paddle.metric.Metric): self.results.append(value.item()) def accumulate(self): + if paddle.distributed.get_world_size() > 1: + results = paddle.to_tensor(self.results) + results_list = [] + paddle.distributed.all_gather(results_list, results) + self.results = paddle.concat(results_list).numpy() + if len(self.results) <= 0: return 0. return np.mean(self.results) diff --git a/ppgan/metrics/psnr_ssim.py b/ppgan/metrics/psnr_ssim.py index 5bd9dc8d4571a25ae2231f52b7ebbd1044a6ed9b..72702de0423b322dacfa676f4b66f534580c781b 100644 --- a/ppgan/metrics/psnr_ssim.py +++ b/ppgan/metrics/psnr_ssim.py @@ -43,6 +43,12 @@ class PSNR(paddle.metric.Metric): self.results.append(value) def accumulate(self): + if paddle.distributed.get_world_size() > 1: + results = paddle.to_tensor(self.results) + results_list = [] + paddle.distributed.all_gather(results_list, results) + self.results = paddle.concat(results_list).numpy() + if len(self.results) <= 0: return 0. return np.mean(self.results)