From 0b3cc6ed1cb5c064692087715e374ee34b10a812 Mon Sep 17 00:00:00 2001
From: Yevgeny Shapiro <yevsh@fb.com>
Date: Wed, 18 May 2022 02:38:07 -0700
Subject: [PATCH] D2Go - Add an evaluator that writes the eval loss to
 tensorboard without recalculation

Summary:
Add an evaluator that writes the eval loss to tensorboard without recalculation, it assumes that the output of the network is a (network_out, losses) tuple.
Also add some typehints.

Differential Revision: D36156130

fbshipit-source-id: cf8d52c1eed0123e1278f93f535591d59571dcf2
---
 detectron2/evaluation/loss_evaluator.py | 56 +++++++++++++++++++++++++
 1 file changed, 56 insertions(+)
 create mode 100644 detectron2/evaluation/loss_evaluator.py

diff --git a/detectron2/evaluation/loss_evaluator.py b/detectron2/evaluation/loss_evaluator.py
new file mode 100644
index 0000000000..b60c74465d
--- /dev/null
+++ b/detectron2/evaluation/loss_evaluator.py
@@ -0,0 +1,56 @@
+from typing import Dict
+
+from detectron2.engine import SimpleTrainer
+from detectron2.evaluation import DatasetEvaluator
+
+
+class LossEvaluator(DatasetEvaluator):
+    """
+    Write validation loss to tensorboard.
+    Assumes that the output of the model arc is a tuple (model_out, losses)
+    """
+    VAL_PREFIX = "Validation_"
+    LOSSES_DICT_STR = "losses_dict"
+
+    def __init__(self):
+        self._aggregated_losses = {}
+        self._num_datapoints = 0
+
+    def process(self, inputs, outputs: Dict):
+        """
+        Add calculated loss to aggregation dict, assume that all losses always exist
+        """
+        # in case this evaluator is invoked in inference
+        if not isinstance(outputs, tuple):
+            return
+
+        assert len(outputs) == 2
+
+        losses_dict = outputs[1]
+        if not isinstance(losses_dict, dict):
+            losses_dict = {'total_loss': losses_dict}
+
+        for loss_name, loss_value in losses_dict.items():
+            validation_loss_name = LossEvaluator.VAL_PREFIX + loss_name
+            loss_value_cpu = loss_value.cpu().detach()
+
+            if self._num_datapoints == 0:
+                self._aggregated_losses[validation_loss_name] = loss_value_cpu
+            else:
+                assert validation_loss_name in self._aggregated_losses, (f"Adding new losses on\
+                    the fly is not supported, loss {validation_loss_name} \
+                    appeared on datapoint {self._num_datapoints}")
+                self._aggregated_losses[validation_loss_name] = (
+                    self._num_datapoints * self._aggregated_losses[validation_loss_name]
+                    + loss_value_cpu) / (self._num_datapoints + 1)
+        self._num_datapoints += 1
+
+    def evaluate(self):
+        try:
+            SimpleTrainer.write_metrics(
+                self._aggregated_losses, 0, LossEvaluator.VAL_PREFIX
+            )
+        except AssertionError as e:
+            # If there is no storage context, this method will not work, otherwise, raise exception
+            if not e.args[0].find("EventStorage") >= 0:
+                raise