[Bug Fix] Move the log with "[Processing]" hint earlier and support outputs of which the type is list of list. (#448)

* Move the log with "[Processing]" hint earlier.

* Flatten all the output tensors to a list when the returned outputs contain list of list.

Author: Xreki

graph_net/paddle/test_compiler.py

@@ -4,10 +4,6 @@
 from pathlib import Path
 import sys
 import os
-from dataclasses import dataclass
-from contextlib import contextmanager
-import time
-import math
 import numpy as np
 import random
 import platform
@@ -62,7 +58,7 @@ def get_hardward_name(args):
                     )
                 )
             )
-        except Exception as e:
+        except Exception:
             pass
     elif args.device == "cpu":
         hardware = platform.processor()
@@ -128,7 +124,7 @@ def get_static_model(args, model):
         backend=None,
     )
     static_model.eval()
-    program = static_model.forward.concrete_program.main_program
+    program = static_model.forward.concrete_program.main_program  # noqa
     return static_model
 
 
@@ -225,47 +221,56 @@ def measure_performance(model_call, args, compiler, profile=False):
 
 
 def check_outputs(args, expected_out, compiled_out):
-    if isinstance(expected_out, paddle.Tensor):
-        expected_out = [expected_out]
-    if isinstance(compiled_out, paddle.Tensor):
-        compiled_out = [compiled_out]
-
-    eager_dtypes = [None] * len(expected_out)
-    for i, tensor in enumerate(expected_out):
-        eager_dtypes[i] = (
-            str(tensor.dtype).replace("paddle.", "") if tensor is not None else "None"
-        )
-
-    compiled_dtypes = [None] * len(compiled_out)
-    for i, tensor in enumerate(compiled_out):
-        compiled_dtypes[i] = (
-            str(tensor.dtype).replace("paddle.", "") if tensor is not None else "None"
-        )
-
+    def _flatten_outputs_to_list(outs):
+        flattened_outs = outs
+        if isinstance(outs, paddle.Tensor):
+            flattened_outs = [outs]
+        else:
+            flattened_outs = [
+                x
+                for out in outs
+                for x in (out if isinstance(out, (tuple, list)) else (out,))
+            ]
+        return flattened_outs
+
+    expected_out = _flatten_outputs_to_list(expected_out)
+    compiled_out = _flatten_outputs_to_list(compiled_out)
+
+    def _get_output_dtypes(outs):
+        dtypes = [
+            str(tensor.dtype).replace("paddle.", "")
+            if isinstance(tensor, paddle.Tensor)
+            else None
+            for i, tensor in enumerate(outs)
+        ]
+        return dtypes
+
+    eager_dtypes = _get_output_dtypes(expected_out)
+    compiled_dtypes = _get_output_dtypes(compiled_out)
     type_match = test_compiler_util.check_output_datatype(
         args, eager_dtypes, compiled_dtypes
     )
 
-    eager_shapes = [None] * len(expected_out)
-    for i, tensor in enumerate(expected_out):
-        eager_shapes[i] = tensor.shape if tensor is not None else None
-
-    compiled_shapes = [None] * len(compiled_out)
-    for i, tensor in enumerate(compiled_out):
-        compiled_shapes[i] = tensor.shape if tensor is not None else None
+    def _get_output_shapes(outs):
+        shapes = [
+            tensor.shape if isinstance(tensor, paddle.Tensor) else None
+            for i, tensor in enumerate(outs)
+        ]
+        return shapes
 
+    eager_shapes = _get_output_shapes(expected_out)
+    compiled_shapes = _get_output_shapes(compiled_out)
     shape_match = test_compiler_util.check_output_shape(
         args, eager_shapes, compiled_shapes
     )
 
-    def transfer_to_float(origin_outputs):
+    def _transfer_to_float(origin_outputs):
         outputs = []
         for item in origin_outputs:
-            if (
-                item is not None
-                and isinstance(item, paddle.Tensor)
-                and item.dtype not in [paddle.float32, paddle.float64]
-            ):
+            if isinstance(item, paddle.Tensor) and item.dtype not in [
+                paddle.float32,
+                paddle.float64,
+            ]:
                 item = item.astype("float32")
             outputs.append(item)
         return outputs
@@ -278,8 +283,8 @@ def transfer_to_float(origin_outputs):
             cmp_equal_func=get_cmp_equal,
         )
 
-        expected_out_fp32 = transfer_to_float(expected_out)
-        compiled_out_fp32 = transfer_to_float(compiled_out)
+        expected_out_fp32 = _transfer_to_float(expected_out)
+        compiled_out_fp32 = _transfer_to_float(compiled_out)
         test_compiler_util.check_allclose(
             args,
             expected_out_fp32,
@@ -308,11 +313,16 @@ def check_and_print_gpu_utilization(compiler):
 
 
 def test_single_model(args):
+    model_path = os.path.normpath(args.model_path)
+    test_compiler_util.print_with_log_prompt(
+        "[Processing]", model_path, args.log_prompt
+    )
+
     compiler = get_compiler_backend(args)
     check_and_print_gpu_utilization(compiler)
 
-    input_dict = get_input_dict(args.model_path)
-    model = get_model(args.model_path)
+    input_dict = get_input_dict(model_path)
+    model = get_model(model_path)
     model.eval()
 
     test_compiler_util.print_basic_config(
@@ -341,7 +351,7 @@ def test_single_model(args):
     compiled_time_stats = {}
     try:
         print("Run model in compiled mode.", file=sys.stderr, flush=True)
-        input_spec = get_input_spec(args.model_path)
+        input_spec = get_input_spec(model_path)
         compiled_model = compiler(model, input_spec)
         compiled_out, compiled_time_stats = measure_performance(
             lambda: compiled_model(**input_dict), args, compiler, profile=False

graph_net/paddle/test_reference_device.py

@@ -1,14 +1,9 @@
 import argparse
-import importlib.util
 import paddle
-import time
-import numpy as np
-import random
 import os
 from pathlib import Path
 from contextlib import redirect_stdout, redirect_stderr
 import json
-import re
 import sys
 import traceback
 
@@ -28,18 +23,23 @@ def get_reference_output_path(reference_dir, model_path):
 
 
 def test_single_model(args):
-    ref_log = get_reference_log_path(args.reference_dir, args.model_path)
-    ref_dump = get_reference_output_path(args.reference_dir, args.model_path)
+    model_path = os.path.normpath(args.model_path)
+    ref_log = get_reference_log_path(args.reference_dir, model_path)
+    ref_dump = get_reference_output_path(args.reference_dir, model_path)
     print(f"Reference log path: {ref_log}", file=sys.stderr, flush=True)
     print(f"Reference outputs path: {ref_dump}", file=sys.stderr, flush=True)
 
     with open(ref_log, "w", encoding="utf-8") as log_f:
         with redirect_stdout(log_f), redirect_stderr(log_f):
+            test_compiler_util.print_with_log_prompt(
+                "[Processing]", model_path, args.log_prompt
+            )
+
             compiler = test_compiler.get_compiler_backend(args)
             test_compiler.check_and_print_gpu_utilization(compiler)
 
-            input_dict = test_compiler.get_input_dict(args.model_path)
-            model = test_compiler.get_model(args.model_path)
+            input_dict = test_compiler.get_input_dict(model_path)
+            model = test_compiler.get_model(model_path)
             model.eval()
 
             test_compiler_util.print_with_log_prompt(
@@ -55,7 +55,7 @@ def test_single_model(args):
             success = False
             time_stats = {}
             try:
-                input_spec = test_compiler.get_input_spec(args.model_path)
+                input_spec = test_compiler.get_input_spec(model_path)
                 compiled_model = compiler(model, input_spec)
                 outputs, time_stats = test_compiler.measure_performance(
                     lambda: compiled_model(**input_dict),

graph_net/paddle/test_target_device.py

@@ -1,12 +1,6 @@
 import argparse
-import importlib.util
-import time
-import numpy as np
-import random
 import os
-from pathlib import Path
 import json
-import re
 import sys
 import traceback
 
@@ -64,11 +58,17 @@ def update_args_and_set_seed(args, model_path):
 
 
 def test_single_model(args):
+    model_path = os.path.normpath(args.model_path)
+    test_compiler_util.print_with_log_prompt(
+        "[Processing]", model_path, args.log_prompt
+    )
+    args = update_args_and_set_seed(args, model_path)
+
     compiler = test_compiler.get_compiler_backend(args)
     test_compiler.check_and_print_gpu_utilization(compiler)
 
-    input_dict = test_compiler.get_input_dict(args.model_path)
-    model = test_compiler.get_model(args.model_path)
+    input_dict = test_compiler.get_input_dict(model_path)
+    model = test_compiler.get_model(model_path)
     model.eval()
 
     test_compiler_util.print_basic_config(
@@ -80,7 +80,7 @@ def test_single_model(args):
     success = False
     time_stats = {}
     try:
-        input_spec = test_compiler.get_input_spec(args.model_path)
+        input_spec = test_compiler.get_input_spec(model_path)
         compiled_model = compiler(model, input_spec)
         outputs, time_stats = test_compiler.measure_performance(
             lambda: compiled_model(**input_dict), args, compiler, profile=False
@@ -95,17 +95,17 @@ def test_single_model(args):
 
     test_compiler_util.print_running_status(args, success)
 
-    model_name = test_compiler_util.get_model_name(args.model_path)
-    if test_compiler_util.get_subgraph_tag(args.model_path):
-        model_name += "_" + test_compiler_util.get_subgraph_tag(args.model_path)
+    model_name = test_compiler_util.get_model_name(model_path)
+    if test_compiler_util.get_subgraph_tag(model_path):
+        model_name += "_" + test_compiler_util.get_subgraph_tag(model_path)
 
     ref_dump = test_reference_device.get_reference_output_path(
-        args.reference_dir, args.model_path
+        args.reference_dir, model_path
     )
     ref_out = paddle.load(str(ref_dump))
 
     ref_log = test_reference_device.get_reference_log_path(
-        args.reference_dir, args.model_path
+        args.reference_dir, model_path
     )
     ref_time_stats = parse_time_stats_from_reference_log(ref_log)
 
@@ -170,7 +170,6 @@ def main(args):
     test_compiler.init_env(args)
 
     if path_utils.is_single_model_dir(args.model_path):
-        args = update_args_and_set_seed(args, args.model_path)
         test_single_model(args)
     else:
         test_multi_models(args)

graph_net/test_compiler_util.py

@@ -143,8 +143,6 @@ def print_with_log_prompt(key, value, log_prompt):
 
 def print_basic_config(args, hardware_name, compile_framework_version):
     model_path = os.path.normpath(args.model_path)
-    print_with_log_prompt("[Processing]", model_path, args.log_prompt)
-
     model_name = get_model_name(model_path)
     print_with_log_prompt("[Config] model:", model_name, args.log_prompt)