[Enhancement] Set real_feat to cpu in inception_utils (#1415)

fix s2 configs
open-mmlab · Nov 7, 2022 · 3349986 · 3349986
1 parent b70695d
commit 3349986
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/configs/styleganv2/stylegan2_c2_8xb4_lsun-car-384x512.py b/configs/styleganv2/stylegan2_c2_8xb4_lsun-car-384x512.py
@@ -98,8 +98,9 @@
 metrics = [
     dict(
         type='FrechetInceptionDistance',
-        prefix='FID-Full-50k',
+        prefix='FID-50k',
         fake_nums=50000,
+        real_nums=50000,
         inception_style='StyleGAN',
         sample_model='ema'),
     dict(type='PrecisionAndRecall', fake_nums=50000, prefix='PR-50K'),
@@ -110,7 +111,7 @@
 #     checkpoint=dict(
 #         save_best=['FID-Full-50k/fid', 'IS-50k/is'],
 #         rule=['less', 'greater']))
-default_hooks = dict(checkpoint=dict(save_best='FID-Full-50k/fid'))
+default_hooks = dict(checkpoint=dict(save_best='FID-50k/fid'))
 
 val_evaluator = dict(metrics=metrics)
 test_evaluator = dict(metrics=metrics)
diff --git a/mmedit/evaluation/functional/inception_utils.py b/mmedit/evaluation/functional/inception_utils.py
@@ -418,7 +418,7 @@ def prepare_inception_feat(dataloader: DataLoader,
                 f'same time. But receive \'{mean}\' and \'{std}\' '
                 'respectively.')
 
-        real_feat_ = metric.forward_inception(img)
+        real_feat_ = metric.forward_inception(img).cpu()
         real_feat.append(real_feat_)
 
         if is_main_process():
@@ -443,7 +443,7 @@ def prepare_inception_feat(dataloader: DataLoader,
     if is_main_process():
         inception_state = dict(**args)
         if capture_mean_cov:
-            real_feat = torch.cat(real_feat, dim=0)[:num_items].cpu().numpy()
+            real_feat = torch.cat(real_feat, dim=0)[:num_items].numpy()
             real_mean = np.mean(real_feat, 0)
             real_cov = np.cov(real_feat, rowvar=False)
             inception_state['real_mean'] = real_mean