Repository: pfnet/pytorch-pfn-extras

Issue:
OutputComparer: add API to limit `max_pool` size
It can be convenient when comparing against engines with limitation (e.g., memory size)


Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/utils/comparer.py b/pytorch_pfn_extras/utils/comparer.py
--- a/pytorch_pfn_extras/utils/comparer.py
+++ b/pytorch_pfn_extras/utils/comparer.py
@@ -93,6 +93,7 @@ class OutputsComparer:
     def __init__(
             self, engines, to_compare_keys=None, *,
             compare_fn=_default_comparer,
+            concurrency=None,
     ):
         """A class for comparison of iteration outputs.
 
@@ -105,6 +106,9 @@ def __init__(
                 A set of keys of output dict to compare.
             compare_fn (function):
                 Comparison function. Default is ``get_default_comparer()``.
+            concurrency (int, optional):
+                The upper bound limit on the number of workers that run concurrently.
+                If ``None``, inferred from the size of ``engines``.
 
         Examples:
             >>> trainer_cpu = ppe.engine.create_trainer(
@@ -137,6 +141,8 @@ def __init__(
         self.report_lock = threading.Lock()
         self.compare_fn = compare_fn
         self._finalized = False
+        self._semaphore = threading.Semaphore(
+            len(engines) if concurrency is None else concurrency)
 
     def _assert_incompatible_trigger(self, condition):
         if not condition:
@@ -160,7 +166,9 @@ def report_output(self, handle, batch_idx, outputs):
                 self._assert_incompatible_trigger(not self._finalized)
 
             # Excplicitly synchronize
+            self._semaphore.release()
             self.barrier.wait()
+            self._semaphore.acquire()
 
     def _compare_outs(self):
         names = list(self.outputs.keys())
@@ -174,6 +182,7 @@ def _compare_outs(self):
 
     def run_engine(self, engine, loaders):
         try:
+            self._semaphore.acquire()
             if isinstance(loaders, tuple):
                 engine.run(*loaders)
             elif isinstance(loaders, dict):
@@ -186,6 +195,8 @@ def run_engine(self, engine, loaders):
         except Exception:
             self.barrier.abort()
             raise
+        finally:
+            self._semaphore.release()
 
     def compare(self, loaders, n_iters=None):
         """Compares outputs.


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/utils_tests/test_comparer.py b/tests/pytorch_pfn_extras_tests/utils_tests/test_comparer.py
--- a/tests/pytorch_pfn_extras_tests/utils_tests/test_comparer.py
+++ b/tests/pytorch_pfn_extras_tests/utils_tests/test_comparer.py
@@ -165,3 +165,42 @@ def test_comparer_incompat_trigger(engine_fn):
             comp.compare({"cpu": (train_1, eval_1), "gpu": (train_2, eval_2)})
         else:
             comp.compare({"cpu": train_1, "gpu": train_2})
+
+
+@pytest.mark.parametrize("engine_fn", [
+    _get_trainer, _get_evaluator, _get_trainer_with_evaluator])
+def test_compare_concurrency(engine_fn):
+    engine_cpu = engine_fn("cpu", 1.0)
+    engine_gpu = engine_fn("cuda:0", 1.0)
+    comp = ppe.utils.comparer.OutputsComparer(
+        {"cpu": engine_cpu, "gpu": engine_gpu}, "a",
+        concurrency=1,
+    )
+    train_1 = list(torch.ones(10) for _ in range(10))
+    train_2 = list(torch.ones(10) for _ in range(10))
+    if engine_fn is _get_trainer_with_evaluator:
+        eval_1 = list(torch.ones(10) for _ in range(10))
+        eval_2 = list(torch.ones(10) for _ in range(10))
+        comp.compare({"cpu": (train_1, eval_1), "gpu": (train_2, eval_2)})
+    else:
+        comp.compare({"cpu": train_1, "gpu": train_2})
+
+
+@pytest.mark.parametrize("engine_fn", [
+    _get_trainer, _get_evaluator, _get_trainer_with_evaluator])
+def test_compare_concurrency_wrong(engine_fn):
+    engine_cpu = engine_fn("cpu", 1.0)
+    engine_gpu = engine_fn("cuda:0", 0.5)
+    comp = ppe.utils.comparer.OutputsComparer(
+        {"cpu": engine_cpu, "gpu": engine_gpu}, "a",
+        concurrency=1,
+    )
+    train_1 = list(torch.ones(10) for _ in range(10))
+    train_2 = list(torch.ones(10) for _ in range(10))
+    with pytest.raises(AssertionError):
+        if engine_fn is _get_trainer_with_evaluator:
+            eval_1 = list(torch.ones(10) for _ in range(10))
+            eval_2 = list(torch.ones(10) for _ in range(10))
+            comp.compare({"cpu": (train_1, eval_1), "gpu": (train_2, eval_2)})
+        else:
+            comp.compare({"cpu": train_1, "gpu": train_2})