add evaluation result type annotation

Summary: Pull Request resolved: https://github.com/facebookresearch/d2go/pull/343

Reviewed By: miqueljubert

Differential Revision: D38077850

fbshipit-source-id: a79541d899ce2b49a30c7f2a81a616f76321026f

d2go/evaluation/api.py

+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from typing import Dict, TypeVar, Union
+
+T = TypeVar("T")
+
+# "accuracy" in D2Go is defined by a 4-level dictionary in the order of:
+# model_tag -> dataset -> task -> metrics
+AccuracyDict = Dict[str, Dict[str, Dict[str, Dict[str, T]]]]
+
+# "metric" in D2Go is a nested dictionary, which may have arbitrary levels.
+MetricsDict = Union[Dict[str, "MetricsDict"], T]

tools/benchmark_data.py

@@ -7,12 +7,13 @@ Tool for benchmarking data loading
 import logging
 import time
 from dataclasses import dataclass
-from typing import Any, Dict, Type, Union
+from typing import Type, Union
 
 import detectron2.utils.comm as comm
 import numpy as np
 from d2go.config import CfgNode
 from d2go.distributed import get_num_processes_per_machine, launch
+from d2go.evaluation.api import AccuracyDict, MetricsDict
 from d2go.runner import BaseRunner
 from d2go.setup import (
     basic_argument_parser,
@@ -30,9 +31,8 @@ logger = logging.getLogger("d2go.tools.benchmark_data")
 
 @dataclass
 class BenchmarkDataOutput:
-    accuracy: Dict[str, Any]
-    # TODO: support arbitrary levels of dicts
-    metrics: Dict[str, Dict[str, Dict[str, Dict[str, float]]]]
+    accuracy: AccuracyDict[float]
+    metrics: MetricsDict[float]
 
 
 def main(

tools/evaluator.py

@@ -9,11 +9,12 @@ torchscript, caffe2, etc.) using Detectron2Go system (dataloading, evaluation, e
 import logging
 import sys
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Type, Union
+from typing import Optional, Type, Union
 
 import torch
 from d2go.config import CfgNode
 from d2go.distributed import launch
+from d2go.evaluation.api import AccuracyDict, MetricsDict
 from d2go.runner import BaseRunner
 from d2go.setup import (
     basic_argument_parser,
@@ -31,9 +32,8 @@ logger = logging.getLogger("d2go.tools.caffe2_evaluator")
 
 @dataclass
 class EvaluatorOutput:
-    accuracy: Dict[str, Any]
-    # TODO: support arbitrary levels of dicts
-    metrics: Dict[str, Dict[str, Dict[str, Dict[str, float]]]]
+    accuracy: AccuracyDict[float]
+    metrics: MetricsDict[float]
 
 
 def main(

tools/train_net.py

@@ -8,11 +8,12 @@ Detection Training Script.
 import logging
 import sys
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Type, Union
+from typing import Dict, List, Optional, Type, Union
 
 import detectron2.utils.comm as comm
 from d2go.config import CfgNode
 from d2go.distributed import launch
+from d2go.evaluation.api import AccuracyDict, MetricsDict
 from d2go.runner import BaseRunner
 from d2go.setup import (
     basic_argument_parser,
@@ -34,9 +35,8 @@ logger = logging.getLogger("d2go.tools.train_net")
 
 @dataclass
 class TrainNetOutput:
-    accuracy: Dict[str, Any]
-    # TODO: support arbitrary levels of dicts
-    metrics: Dict[str, Dict[str, Dict[str, Dict[str, Any]]]]
+    accuracy: AccuracyDict[float]
+    metrics: MetricsDict[float]
     model_configs: Dict[str, str]
     # TODO: decide if `tensorboard_log_dir` should be part of output
     tensorboard_log_dir: Optional[str] = None