Define a unified BenchmarkResult and support write to json.

Author: Xreki

graph_net/benchmark_result.py

@@ -0,0 +1,73 @@
+import os
+import sys
+import json
+import re
+
+
+class BenchmarkResult:
+    def __init__(self, args, hardware, compile_framework_version):
+        self.configuration = {
+            "model_name": self.get_model_name(args),
+            "device": args.device,
+            "hardware": hardware,
+            "compiler": args.compiler,
+            "compile_framework_version": compile_framework_version,
+            "warmup": args.warmup,
+            "trials": args.trials,
+        }
+        self.model_info = {
+            "num_ops": -1,
+            "input_dtypes": None,
+            "param_dtypes": None,
+        }
+        self.correctness = {}
+        self.performance = {
+            "eager": None,
+            "compiled": None,
+            "speedup": None,
+        }
+
+    def get_model_name(self, args):
+        fields = args.model_path.split(os.sep)
+
+        pattern = rf"^subgraph(_\d+)?$"
+        if re.match(pattern, fields[-1]):
+            model_name = f"{fields[-2]}_{fields[-1]}"
+        else:
+            model_name = fields[-1]
+        return model_name
+
+    def update_model_info(self, num_ops, input_dtypes, param_dtypes):
+        self.model_info["num_ops"] = num_ops
+        self.model_info["input_dtypes"] = input_dtypes
+        self.model_info["param_dtypes"] = param_dtypes
+
+    def update_corrrectness(self, key, cmp_ret):
+        self.correctness[key] = cmp_ret
+
+    def update_performance(self, eager_time_ms, compiled_time_ms):
+        self.performance["eager"] = eager_time_ms
+        self.performance["compiled"] = compiled_time_ms
+        if eager_time_ms > 0 and compiled_time_ms > 0:
+            self.performance["speedup"] = eager_time_ms / compiled_time_ms
+        return self.performance["speedup"]
+
+    def write_to_json(self, output_dir):
+        assert output_dir is not None
+        os.makedirs(output_dir, exist_ok=True)
+        result_data = {
+            "configuration": self.configuration,
+            "model_info": self.model_info,
+            "correctness": self.correctness,
+            "performance": {
+                k: float(f"{v:.6f}") if isinstance(v, float) else v
+                for k, v in self.performance.items()
+            },
+        }
+        model_name = self.configuration["model_name"]
+        compiler_name = self.configuration["compiler"]
+        file_path = os.path.join(output_dir, f"{model_name}_{compiler_name}.json")
+        with open(file_path, "w") as f:
+            json.dump(result_data, f, indent=4)
+        print(f"Result saved to {file_path}", file=sys.stderr)
+        print(result_data)

graph_net/paddle/test_compiler.py

@@ -9,8 +9,10 @@
 import time
 import numpy as np
 import random
+import platform
 
-from . import utils
+from graph_net.paddle import utils
+from graph_net.benchmark_result import BenchmarkResult
 
 
 def load_class_from_file(file_path: str, class_name: str):
@@ -201,6 +203,27 @@ def measure_performance(model_call, synchronizer_func, args, profile=False):
     return outs, times
 
 
+def init_benchmark_result(args):
+    if args.device == "cuda":
+        hardware = paddle.device.cuda.get_device_name(0)
+    elif args.device == "cpu":
+        hardware = platform.processor()
+    else:
+        hardware = "unknown"
+
+    if args.compiler == "CINN":
+        compile_framework_version = paddle.__version__
+    else:
+        compile_framework_version = "unknown"
+
+    result_data = BenchmarkResult(
+        args=args,
+        hardware=hardware,
+        compile_framework_version=compile_framework_version,
+    )
+    return result_data
+
+
 def test_single_model(args):
     synchronizer_func = get_synchronizer_func(args)
     input_dict, input_dtypes, param_dtypes = get_input_dict(args)
@@ -210,12 +233,16 @@ def test_single_model(args):
     # Collect model information
     num_ops = count_number_of_ops(args, model)
 
-    print("Run on eager mode")
+    # Initialize benchmark result
+    result_data = init_benchmark_result(args)
+    result_data.update_model_info(num_ops, input_dtypes, param_dtypes)
+
+    # Run on eager mode
     expected_out, eager_time_ms = measure_performance(
         lambda: model(**input_dict), synchronizer_func, args, profile=False
     )
 
-    print("Run on compiling mode")
+    # Run on compiling mode
     compiled_model = get_compiled_model(args, model)
     compiled_out, compiled_time_ms = measure_performance(
         lambda: compiled_model(**input_dict), synchronizer_func, args, profile=False
@@ -243,6 +270,7 @@ def test_single_model(args):
 
     def print_cmp(key, func, **kwargs):
         cmp_ret = func(expected_out, compiled_out, **kwargs)
+        result_data.update_corrrectness(key, cmp_ret)
         print(
             f"{args.log_prompt} {key} model_path:{args.model_path} {cmp_ret}",
             file=sys.stderr,
@@ -271,6 +299,10 @@ def print_cmp(key, func, **kwargs):
         file=sys.stderr,
     )
 
+    result_data.update_performance(eager_time_ms, compiled_time_ms)
+    if args.output_dir:
+        result_data.write_to_json(args.output_dir)
+
 
 def get_cmp_equal(expected_out, compiled_out):
     return " ".join(
@@ -372,6 +404,13 @@ def main(args):
         default="CINN",
         help="Path to customized compiler python file",
     )
+    parser.add_argument(
+        "--device",
+        type=str,
+        required=False,
+        default="cuda",
+        help="Device for testing the compiler (e.g., 'cpu' or 'cuda')",
+    )
     parser.add_argument(
         "--warmup", type=int, required=False, default=5, help="Number of warmup steps"
     )
@@ -391,5 +430,12 @@ def main(args):
         default="graph-net-test-compiler-log",
         help="Log prompt for performance log filtering.",
     )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        required=False,
+        default=None,
+        help="Directory to save the structured JSON result file.",
+    )
     args = parser.parse_args()
     main(args=args)