Implement FlowGaussianNll evaluator

codes/models/eval/__init__.py

@@ -1,3 +1,4 @@
+from models.eval.flow_gaussian_nll import FlowGaussianNll
 from models.eval.sr_style import SrStyleTransferEvaluator
 from models.eval.style import StyleTransferEvaluator
 
@@ -8,5 +9,7 @@ def create_evaluator(model, opt_eval, env):
         return StyleTransferEvaluator(model, opt_eval, env)
     elif type == 'sr_stylegan':
         return SrStyleTransferEvaluator(model, opt_eval, env)
+    elif type == 'flownet_gaussian':
+        return FlowGaussianNll(model, opt_eval, env)
     else:
         raise NotImplementedError()

codes/models/eval/flow_gaussian_nll.py

@@ -0,0 +1,37 @@
+import os
+
+import torch
+import os.path as osp
+import torchvision
+from torch.utils.data import DataLoader
+
+import models.eval.evaluator as evaluator
+from pytorch_fid import fid_score
+
+
+# Evaluate how close to true Gaussian a flow network predicts in a "normal" pass given a LQ/HQ image pair.
+from data.image_folder_dataset import ImageFolderDataset
+from models.archs.srflow_orig.flow import GaussianDiag
+
+
+class FlowGaussianNll(evaluator.Evaluator):
+    def __init__(self, model, opt_eval, env):
+        super().__init__(model, opt_eval, env)
+        self.batch_sz = opt_eval['batch_size']
+        self.dataset = ImageFolderDataset(opt_eval['dataset'])
+        self.dataloader = DataLoader(self.dataset, self.batch_sz)
+
+    def perform_eval(self):
+        total_zs = 0
+        z_loss = 0
+        with torch.no_grad():
+            for batch in self.dataloader:
+                z, _, _ = self.model(gt=batch['GT'],
+                                     lr=batch['LQ'],
+                                     epses=[],
+                                     reverse=False,
+                                     add_gt_noise=False)
+                for z_ in z:
+                    z_loss += GaussianDiag.logp(None, None, z_).mean()
+                    total_zs += 1
+        return {"gaussian_diff": z_loss / total_zs}

codes/train.py

@@ -291,7 +291,7 @@ class Trainer:
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_exd_imgsetext_rrdb4x_6bl_bigbatch.yml')
+    parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_exd_imgsetext_rrdb_bigboi.yml')
     parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher')
     parser.add_argument('--local_rank', type=int, default=0)
     args = parser.parse_args()