From 8a00f15746ff094b36e6053d0cbc8bcd872c385a Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Wed, 2 Dec 2020 14:09:54 -0700
Subject: [PATCH] Implement FlowGaussianNll evaluator

---
 codes/models/eval/__init__.py          |  3 +++
 codes/models/eval/flow_gaussian_nll.py | 37 ++++++++++++++++++++++++++
 codes/train.py                         |  2 +-
 3 files changed, 41 insertions(+), 1 deletion(-)
 create mode 100644 codes/models/eval/flow_gaussian_nll.py

diff --git a/codes/models/eval/__init__.py b/codes/models/eval/__init__.py
index 111ca71f..1a612f32 100644
--- a/codes/models/eval/__init__.py
+++ b/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()
\ No newline at end of file
diff --git a/codes/models/eval/flow_gaussian_nll.py b/codes/models/eval/flow_gaussian_nll.py
new file mode 100644
index 00000000..2405a70d
--- /dev/null
+++ b/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}
diff --git a/codes/train.py b/codes/train.py
index 36fae2c5..9995b9f6 100644
--- a/codes/train.py
+++ b/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()