From ae8de0cb9d083b04bb2286decc7e8743a40e54e0 Mon Sep 17 00:00:00 2001 From: James Betker Date: Tue, 15 Jun 2021 10:31:07 -0600 Subject: [PATCH] fid saving images across all rank fix --- codes/trainer/eval/sr_diffusion_fid.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/codes/trainer/eval/sr_diffusion_fid.py b/codes/trainer/eval/sr_diffusion_fid.py index eba49cd7..1891487f 100644 --- a/codes/trainer/eval/sr_diffusion_fid.py +++ b/codes/trainer/eval/sr_diffusion_fid.py @@ -50,9 +50,10 @@ class SrDiffusionFidEvaluator(evaluator.Evaluator): torch.distributed.all_gather(gather_list, gen) gen = torch.cat(gather_list, dim=0) - for g in gen: - torchvision.utils.save_image(g, osp.join(fid_fake_path, f"{counter}.png")) - counter += 1 + if self.env['rank'] <= 0: + for g in gen: + torchvision.utils.save_image(g, osp.join(fid_fake_path, f"{counter}.png")) + counter += 1 return {"fid": fid_score.calculate_fid_given_paths([self.fid_real_samples, fid_fake_path], self.fid_batch_size, True, 2048)} \ No newline at end of file