[torch-nightly] Add v3 score field

userbenchmark/torch-nightly/__init__.py

@@ -3,11 +3,13 @@
 """
 import argparse
 import itertools
+import json
+import math
 import os
 import yaml
 import numpy
 
-from typing import List, Tuple, Dict, Optional
+from typing import List, Tuple, Dict, Optional, Any
 from ..utils import REPO_PATH, add_path, get_output_json, dump_output
 
 with add_path(REPO_PATH):
@@ -17,6 +19,9 @@
 
 BM_NAME = "torch-nightly"
 CURRENT_DIR = os.path.dirname(os.path.realpath(__file__))
+DEFAULT_DELTA_THRESHOLD = 0.07
+DEFAULT_TARGET_SCORE = 1000.0
+
 
 def generate_model_configs(devices: List[str], tests: List[str], model_names: List[str]) -> List[TorchBenchModelConfig]:
     """Use the default batch size and default mode."""
@@ -37,13 +42,36 @@ def generate_model_configs(devices: List[str], tests: List[str], model_names: Li
 def get_metrics(_config: TorchBenchModelConfig) -> List[str]:
     return ["latencies", "cpu_peak_mem", "gpu_peak_mem"]
 
+def compute_score(results, reference_latencies: Dict[str, float]) -> float:
+    # sanity checks
+    latency_results = {k: v for k, v in results.items() if k.endswith("_latency")}
+    test_set = set(latency_results.keys())
+    reference_set = set(reference_latencies.keys())
+    test_only_set = test_set.difference(reference_set)
+    assert not test_only_set, f"Tests {test_only_set} only appears in result json, not in reference yaml."
+    reference_only_set = reference_set.difference(test_set)
+    assert not reference_only_set, f"Tests {reference_only_set} only appears in reference yaml, not in result json."
+    # check that for every test in reference_latencies, we can find the corresponding tests in latency_results
+    total_score = 0.0
+    weight = 1.0 / len(reference_latencies)
+    for key, ref_latency in reference_latencies.items():
+        test_latency = latency_results[key]
+        ref_latency = float(ref_latency)
+        delta = (test_latency - ref_latency) / test_latency
+        # If less than threshold, treat it as noise
+        if abs(delta) <= DEFAULT_DELTA_THRESHOLD:
+            test_latency = ref_latency
+        total_score += weight * math.log(ref_latency / test_latency)
+    score = math.exp(total_score) * DEFAULT_TARGET_SCORE
+    return score
+
 def result_to_output_metrics(results: List[Tuple[TorchBenchModelConfig, TorchBenchModelMetrics]]) -> Dict[str, float]:
     # metrics name examples:
     # test_eval[timm_regnet-cuda-eager]_latency
     # test_eval[timm_regnet-cuda-eager]_cmem
     # test_eval[timm_regnet-cuda-eager]_gmem
     result_metrics = {}
-    for config, metrics in results:
+    for config_id, (config, metrics) in enumerate(results):
         metrics_base = f"test_{config.test}[{config.name}-{config.device}-eager]"
         latency_metric = f"{metrics_base}_latency"
         median_latency = numpy.median(metrics.latencies)
@@ -67,12 +95,13 @@ def validate(candidates: List[str], choices: List[str]) -> List[str]:
         assert candidate in choices, f"Specified {candidate}, but not in available list: {choices}."
     return candidates
 
-def generate_model_configs_from_yaml(yaml_file: str) -> List[TorchBenchModelConfig]:
+def generate_model_configs_from_yaml(yaml_file: str) -> Tuple[TorchBenchModelConfig, List[float], Any]:
     yaml_file_path = os.path.join(CURRENT_DIR, yaml_file)
     with open(yaml_file_path, "r") as yf:
         config_obj = yaml.safe_load(yf)
-    devices = config_obj.keys()
+    devices = config_obj["metadata"]["devices"]
     configs = []
+    reference_latencies = {}
     for device in devices:
         for c in config_obj[device]:
             if not c["stable"]:
@@ -87,7 +116,10 @@ def generate_model_configs_from_yaml(yaml_file: str) -> List[TorchBenchModelConf
                 extra_env=None,
             )
             configs.append(config)
-    return configs
+            metrics_base = f"test_{config.test}[{config.name}-{config.device}-eager]"
+            latency_metric_key = f"{metrics_base}_latency"
+            reference_latencies[latency_metric_key] = c["median_latency"]
+    return configs, reference_latencies, config_obj
 
 def parse_str_to_list(candidates):
     if isinstance(candidates, list):
@@ -120,12 +152,24 @@ def parse_args(args):
     parser.add_argument("--model", "-m", default=None, type=str, help="Only run the specifice models, splited by comma.")
     parser.add_argument("--config", "-c", default=None, help="YAML config to specify tests to run.")
     parser.add_argument("--dryrun", action="store_true", help="Dryrun the command.")
+    parser.add_argument("--score", default=None, help="Generate score from the past run json.")
     return parser.parse_args(args)
 
 def run(args: List[str]):
     args = parse_args(args)
-    if args.config:
-        configs = generate_model_configs_from_yaml(args.config)
+    if args.score:
+        assert args.config, f"To compute score, you must specify the config YAML using --config."
+        configs, reference_latencies, config_obj = generate_model_configs_from_yaml(args.config)
+        with open(args.score, "r") as sp:
+            run_result = json.load(sp)
+        input_metrics = run_result["metrics"]
+        score = compute_score(input_metrics, reference_latencies)
+        score_version = config_obj["metadata"]["score_version"]
+        score_name = f"{score_version}_score"
+        print(f"TorchBench {score_name}: {score}.")
+        exit(0)
+    elif args.config:
+        configs, reference_latencies = generate_model_configs_from_yaml(args.config)
     else:
         # If not specified, use the entire model set
         if not args.model:
@@ -134,6 +178,7 @@ def run(args: List[str]):
         tests = validate(parse_str_to_list(args.test), list_tests())
         models = validate(parse_str_to_list(args.model), list_models())
         configs = generate_model_configs(devices, tests, model_names=models)
+        reference_latencies = None
     results = []
     try:
         for config in configs:
@@ -144,4 +189,9 @@ def run(args: List[str]):
         print("User keyboard interrupted!")
     if not args.dryrun:
         metrics = result_to_output_metrics(results)
+        if reference_latencies:
+            score = compute_score(metrics, reference_latencies)
+            score_version = config_obj["metadata"]["score_version"]
+            score_name = f"{score_version}_score"
+            metrics[score_name] = score
         dump_result_to_json(metrics)