Merge pull request #195 from argoai/limit-det-eval-procs

Require manual specification of number of processes to use in detection eval

argoverse/evaluation/detection/eval.py

@@ -64,7 +64,7 @@
 from collections import defaultdict
 from multiprocessing import Pool
 from pathlib import Path
-from typing import DefaultDict, List, NamedTuple
+from typing import DefaultDict, List
 
 import numpy as np
 import pandas as pd
@@ -76,25 +76,37 @@
 logger = logging.getLogger(__name__)
 
 
-class DetectionEvaluator(NamedTuple):
-    """Instantiates a DetectionEvaluator object for evaluation.
+class DetectionEvaluator:
+    """Instantiates a DetectionEvaluator object for evaluation."""
 
-    Args:
-        dt_fpath_root: Path to the folder which contains the detections.
-        gt_fpath_root: Path to the folder which contains the split of logs.
-        figs_fpath: Path to the folder which will contain the output figures.
-        cfg: Detection configuration settings.
-    """
-
-    dt_root_fpath: Path
-    gt_root_fpath: Path
-    figs_fpath: Path
-    cfg: DetectionCfg = DetectionCfg()
+    def __init__(
+        self,
+        dt_root_fpath: Path,
+        gt_root_fpath: Path,
+        figs_fpath: Path,
+        cfg: DetectionCfg = DetectionCfg(),
+        num_procs: int = -1,
+    ) -> None:
+        """
+        Args:
+            dt_fpath_root: Path to the folder which contains the detections.
+            gt_fpath_root: Path to the folder which contains the split of logs.
+            figs_fpath: Path to the folder which will contain the output figures.
+            cfg: Detection configuration settings.
+            num_procs: Number of processes among which to subdivide work.
+                Specifying -1 will use one process per available core
+        """
+        self.dt_root_fpath = dt_root_fpath
+        self.gt_root_fpath = gt_root_fpath
+        self.figs_fpath = figs_fpath
+        self.cfg = cfg
+        self.num_procs = os.cpu_count() if num_procs == -1 else num_procs
 
     def evaluate(self) -> pd.DataFrame:
         """Evaluate detection output and return metrics. The multiprocessing
-        library is used for parallel assignment between detections and ground truth
-        annotations.
+        library is used for parallel processing of sweeps -- each sweep is
+        processed independently, computing assignment between detections and
+        ground truth annotations.
 
         Returns:
             Evaluation metrics of shape (C + 1, K) where C + 1 is the number of classes.
@@ -111,7 +123,7 @@ def evaluate(self) -> pd.DataFrame:
         cls_to_ninst: DefaultDict[str, int] = defaultdict(int)
 
         args = [(self.dt_root_fpath, gt_fpath, self.cfg, avm) for gt_fpath in gt_fpaths]
-        with Pool(os.cpu_count()) as p:
+        with Pool(self.num_procs) as p:
             accum = p.starmap(accumulate, args)
 
         for frame_stats, frame_cls_to_inst in accum:
@@ -204,6 +216,14 @@ def main() -> None:
         help="Ground truth root folder path.",
         required=True,
     )
+    parser.add_argument(
+        "-p",
+        "--num_processes",
+        type=int,
+        help="Number of processes among which to subdivide work. Specifying -1 will use one process per available core",
+        default=-1,
+    )
+
     parser.add_argument("-f", "--fig_fpath", type=str, help="Figures root folder path.", default="figs")
     args = parser.parse_args()
     logger.info(f"args == {args}")
@@ -212,7 +232,7 @@ def main() -> None:
     gt_fpath = Path(args.gt_fpath)
     fig_fpath = Path(args.fig_fpath)
 
-    evaluator = DetectionEvaluator(dt_fpath, gt_fpath, fig_fpath)
+    evaluator = DetectionEvaluator(dt_fpath, gt_fpath, fig_fpath, num_procs=args.num_processes)
     metrics = evaluator.evaluate()
     print(metrics)