add default_writers()

Summary: Allow this logic to be reused when not using DefaultTrainer

Reviewed By: theschnitz

Differential Revision: D26213772

fbshipit-source-id: 9c85412fb24323018143c9fe9d42e23562dbb2ff

Author: ppwwyyxx

detectron2/engine/defaults.py

@@ -14,6 +14,7 @@
 import os
 import sys
 from collections import OrderedDict
+from typing import Optional
 import torch
 from fvcore.nn.precise_bn import get_bn_modules
 from torch.nn.parallel import DistributedDataParallel
@@ -43,7 +44,13 @@
 from . import hooks
 from .train_loop import AMPTrainer, SimpleTrainer, TrainerBase
 
-__all__ = ["default_argument_parser", "default_setup", "DefaultPredictor", "DefaultTrainer"]
+__all__ = [
+    "default_argument_parser",
+    "default_setup",
+    "default_writers",
+    "DefaultPredictor",
+    "DefaultTrainer",
+]
 
 
 def default_argument_parser(epilog=None):
@@ -157,6 +164,27 @@ def default_setup(cfg, args):
         torch.backends.cudnn.benchmark = cfg.CUDNN_BENCHMARK
 
 
+def default_writers(output_dir: str, max_iter: Optional[int] = None):
+    """
+    Build a list of :class:`EventWriter` to be used.
+    It now consists of a :class:`CommonMetricPrinter`,
+    :class:`TensorboardXWriter` and :class:`JSONWriter`.
+
+    Args:
+        output_dir: directory to store JSON metrics and tensorboard events
+        max_iter: the total number of iterations
+
+    Returns:
+        list[EventWriter]: a list of :class:`EventWriter` objects.
+    """
+    return [
+        # It may not always print what you want to see, since it prints "common" metrics only.
+        CommonMetricPrinter(max_iter),
+        JSONWriter(os.path.join(output_dir, "metrics.json")),
+        TensorboardXWriter(output_dir),
+    ]
+
+
 class DefaultPredictor:
     """
     Create a simple end-to-end predictor with the given config that runs on
@@ -377,37 +405,21 @@ def test_and_save_results():
         ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
 
         if comm.is_main_process():
+            # Here the default print/log frequency of each writer is used.
             # run writers in the end, so that evaluation metrics are written
             ret.append(hooks.PeriodicWriter(self.build_writers(), period=20))
         return ret
 
     def build_writers(self):
         """
-        Build a list of writers to be used. By default it contains
-        writers that write metrics to the screen,
-        a json file, and a tensorboard event file respectively.
+        Build a list of writers to be used using :func:`default_writers()`.
         If you'd like a different list of writers, you can overwrite it in
         your trainer.
 
         Returns:
             list[EventWriter]: a list of :class:`EventWriter` objects.
-
-        It is now implemented by:
-        ::
-            return [
-                CommonMetricPrinter(self.max_iter),
-                JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
-                TensorboardXWriter(self.cfg.OUTPUT_DIR),
-            ]
-
         """
-        # Here the default print/log frequency of each writer is used.
-        return [
-            # It may not always print what you want to see, since it prints "common" metrics only.
-            CommonMetricPrinter(self.max_iter),
-            JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
-            TensorboardXWriter(self.cfg.OUTPUT_DIR),
-        ]
+        return default_writers(self.cfg.OUTPUT_DIR, self.max_iter)
 
     def train(self):
         """

tests/test_engine.py

@@ -5,7 +5,6 @@
 import tempfile
 import time
 import unittest
-from unittest.mock import MagicMock
 import torch
 from torch import nn
 
@@ -61,12 +60,12 @@ def test_writer_hooks(self):
         with tempfile.TemporaryDirectory(prefix="detectron2_test") as d:
             json_file = os.path.join(d, "metrics.json")
             writers = [CommonMetricPrinter(max_iter), JSONWriter(json_file)]
-            logger_info = writers[0].logger.info = MagicMock()
 
             trainer.register_hooks(
                 [hooks.EvalHook(0, lambda: {"metric": 100}), hooks.PeriodicWriter(writers)]
             )
-            trainer.train(0, max_iter)
+            with self.assertLogs(writers[0].logger) as logs:
+                trainer.train(0, max_iter)
 
             with open(json_file, "r") as f:
                 data = [json.loads(line.strip()) for line in f]
@@ -75,12 +74,11 @@ def test_writer_hooks(self):
                 self.assertIn("metric", data[-1], "Eval metric must be in last line of JSON!")
 
             # test logged messages from CommonMetricPrinter
-            all_logs = [str(x) for x in logger_info.call_args_list]
-            self.assertEqual(len(all_logs), 3)
-            for log, iter in zip(all_logs, [19, 39, 49]):
+            self.assertEqual(len(logs.output), 3)
+            for log, iter in zip(logs.output, [19, 39, 49]):
                 self.assertIn(f"iter: {iter}", log)
 
-            self.assertIn("eta: 0:00:00", all_logs[-1], "Last ETA must be 0!")
+            self.assertIn("eta: 0:00:00", logs.output[-1], "Last ETA must be 0!")
 
     @unittest.skipIf(os.environ.get("CI"), "Require COCO data.")
     def test_default_trainer(self):

tests/test_events.py

@@ -57,8 +57,8 @@ def testPrintETA(self):
 
             with self.assertLogs("detectron2.utils.events") as logs:
                 p1.write()
-            assert "eta" in logs.output[0]
+            self.assertIn("eta", logs.output[0])
 
             with self.assertLogs("detectron2.utils.events") as logs:
                 p2.write()
-            assert "eta" not in logs.output[0]
+            self.assertNotIn("eta", logs.output[0])

tools/plain_train_net.py

@@ -33,7 +33,7 @@
     build_detection_test_loader,
     build_detection_train_loader,
 )
-from detectron2.engine import default_argument_parser, default_setup, launch
+from detectron2.engine import default_argument_parser, default_setup, default_writers, launch
 from detectron2.evaluation import (
     CityscapesInstanceEvaluator,
     CityscapesSemSegEvaluator,
@@ -48,12 +48,7 @@
 )
 from detectron2.modeling import build_model
 from detectron2.solver import build_lr_scheduler, build_optimizer
-from detectron2.utils.events import (
-    CommonMetricPrinter,
-    EventStorage,
-    JSONWriter,
-    TensorboardXWriter,
-)
+from detectron2.utils.events import EventStorage
 
 logger = logging.getLogger("detectron2")
 
@@ -138,15 +133,7 @@ def do_train(cfg, model, resume=False):
         checkpointer, cfg.SOLVER.CHECKPOINT_PERIOD, max_iter=max_iter
     )
 
-    writers = (
-        [
-            CommonMetricPrinter(max_iter),
-            JSONWriter(os.path.join(cfg.OUTPUT_DIR, "metrics.json")),
-            TensorboardXWriter(cfg.OUTPUT_DIR),
-        ]
-        if comm.is_main_process()
-        else []
-    )
+    writers = default_writers(cfg.OUTPUT_DIR, max_iter) if comm.is_main_process() else []
 
     # compared to "train_net.py", we do not support accurate timing and
     # precise BN here, because they are not trivial to implement in a small training loop