Merge branch 'collect_info' into add_cv_samples_5_need_fix

Author: Xreki

graph_net/collect_stats_util.py

@@ -1,4 +1,5 @@
 import ast
+import json
 import importlib
 import inspect
 from dataclasses import dataclass, field
@@ -28,6 +29,7 @@ class ModelStats:
     model_size_in_billion: float = None
     input_dtypes: Dict[str, int] = field(default_factory=dict)
     param_dtypes: Dict[str, int] = field(default_factory=dict)
+    input_shapes: Dict[str, list] = field(default_factory=dict)
     op_dtypes: Dict[str, int] = field(default_factory=dict)
     ops: Dict[str, int] = field(default_factory=dict)
     source: str = None
@@ -37,10 +39,6 @@ class ModelStats:
 def print_model_stats(stats, log_prompt):
     assert isinstance(stats, ModelStats), f"{type(stats)=}"
 
-    def dict_to_string(d):
-        kv_list = [f"{k}:{v}" for k, v in d.items()]
-        return " ".join(kv_list)
-
     def print_with_log_prompt(key, value):
         print(
             f"{log_prompt} [ModelStats.{key}] model_path:{stats.model_path} {value}",
@@ -52,10 +50,11 @@ def print_with_log_prompt(key, value):
     print_with_log_prompt("num_outputs", stats.num_outputs)
     print_with_log_prompt("num_ops", stats.num_ops)
     print_with_log_prompt("model_size", f"{stats.model_size_in_billion}B")
-    print_with_log_prompt("input_dtypes", dict_to_string(stats.input_dtypes))
-    print_with_log_prompt("param_dtypes", dict_to_string(stats.param_dtypes))
-    print_with_log_prompt("op_dtypes", dict_to_string(stats.op_dtypes))
-    print_with_log_prompt("ops", dict_to_string(stats.ops))
+    print_with_log_prompt("input_dtypes", json.dumps(stats.input_dtypes))
+    print_with_log_prompt("param_dtypes", json.dumps(stats.param_dtypes))
+    print_with_log_prompt("input_shapes", json.dumps(stats.input_shapes))
+    print_with_log_prompt("op_dtypes", json.dumps(stats.op_dtypes))
+    print_with_log_prompt("ops", json.dumps(stats.ops))
     print_with_log_prompt("source", stats.source)
     print_with_log_prompt("heuristic_tag", stats.heuristic_tag)
 

graph_net/paddle/collect_stats.py

@@ -109,10 +109,11 @@ def __call__(self, program):
                             else:
                                 # for paddle.base.libpaddle.pir.VectorType, but cannot be accurately determined
                                 if op_name in [
+                                    "broadcast_tensors",
+                                    "distribute_fpn_proposals",
+                                    "meshgrid",
                                     "split",
                                     "split_with_num",
-                                    "meshgrid",
-                                    "distribute_fpn_proposals",
                                 ]:
                                     op_dtype = self.parse_pir_value_dtypes(
                                         str(out.type())
@@ -165,6 +166,7 @@ def collect_model_stats(model_path, log_prompt):
     model = model_class()
 
     model_size = 0
+    input_shapes = set()
     input_dtypes = {}
     param_dtypes = {}
     ops_count_dict = {}
@@ -190,6 +192,7 @@ def collect_model_stats(model_path, log_prompt):
                 param_dtypes[dtype_str] = param_dtypes.get(dtype_str, 0) + 1
             elif name in inputs.keys():
                 input_dtypes[dtype_str] = input_dtypes.get(dtype_str, 0) + 1
+                input_shapes.add(str(value["shape"]))
 
     num_outputs = collect_stats_util.get_number_of_returns(
         file_path, "GraphModule", "forward"
@@ -200,7 +203,7 @@ def collect_model_stats(model_path, log_prompt):
         program_analyzer.is_complete if program_analyzer is not None else False
     )
     print(
-        f"model_stats collection information: model_path={model_path}, method=to_static, is_ops_complete={is_complete}"
+        f"model_stats collection information: model_path={model_path} method=to_static is_ops_complete={is_complete}"
     )
 
     stats = collect_stats_util.ModelStats(
@@ -212,6 +215,7 @@ def collect_model_stats(model_path, log_prompt):
         model_size_in_billion=model_size / 1e9,
         input_dtypes=input_dtypes,
         param_dtypes=param_dtypes,
+        input_shapes=list(input_shapes),
         op_dtypes=op_dtypes,
         ops=ops_count_dict,
         source=source,

graph_net/torch/collect_stats.py

@@ -336,17 +336,23 @@ def collect_model_stats(model_path, device, log_prompt):
                     op_dtypes[dtype_str] = op_dtypes.get(dtype_str, 0) + num
 
     model_size = 0
+    input_shapes = set()
     input_dtypes = {}
     param_dtypes = {}
     for name, arg_type in argument_name2types.items():
-        if arg_type == torch.nn.parameter.Parameter:
+        if (
+            name.startswith("L_self_modules_")
+            or arg_type == torch.nn.parameter.Parameter
+        ):
+            # Some parameters like L_self_modules_bn1_buffers_running_mean_ are torch.Tensor.
             param_numel = math.prod(input_dict[name].shape)
             model_size += param_numel
             dtype_str = str(input_dict[name].dtype).replace("torch.", "")
             param_dtypes[dtype_str] = param_dtypes.get(dtype_str, 0) + 1
-        else:
+        elif arg_type == torch.Tensor:
             dtype_str = str(input_dict[name].dtype).replace("torch.", "")
             input_dtypes[dtype_str] = input_dtypes.get(dtype_str, 0) + 1
+            input_shapes.add(str(list(input_dict[name].shape)))
 
     num_outputs = collect_stats_util.get_number_of_returns(
         file_path, "GraphModule", "forward"
@@ -356,7 +362,7 @@ def collect_model_stats(model_path, device, log_prompt):
 
     is_complete = meta_executor.is_complete if meta_executor is not None else False
     print(
-        f"model_stats collection information: model_path={model_path}, method={method}, is_ops_complete={is_complete}"
+        f"model_stats collection information: model_path={model_path} method={method} is_ops_complete={is_complete}"
     )
 
     stats = collect_stats_util.ModelStats(
@@ -368,6 +374,7 @@ def collect_model_stats(model_path, device, log_prompt):
         model_size_in_billion=model_size / 1e9,
         input_dtypes=input_dtypes,
         param_dtypes=param_dtypes,
+        input_shapes=list(input_shapes),
         op_dtypes=op_dtypes,
         ops=ops_count_dict,
         source=source,