diff --git a/padiff/abstracts/marker.py b/padiff/abstracts/marker.py index b285e4a..7f00a00 100644 --- a/padiff/abstracts/marker.py +++ b/padiff/abstracts/marker.py @@ -62,7 +62,7 @@ def update_unassigned_weights_list(self, layers, mode="all"): def set_layer_map(self, layer_map): _layer_map = [] - for layer in self.traversal_for_assign_weight(): + for layer in self.traversal_for_layer_map(): if layer.model in layer_map: self.unassigned_weights_list_recursively.add(layer.model) _layer_map.append(layer) @@ -79,7 +79,7 @@ def auto_layer_map(self, model_place): registered = init_pool.registered_base_models if model_place == "base" else init_pool.registered_raw_models log("Auto set layer_map start searching...") - for layer in self.traversal_for_auto_layer_map(): + for layer in self.traversal_for_layer_map(): if layer.fullname in registered: print(f"++++ {model_place}_model found `{layer.fullname}` add to layer_map ++++") _layer_map.append(layer) @@ -105,11 +105,14 @@ def traversal_for_hook(self): def traversal_for_assign_weight(self): yield self.proxy_model for model in traversal_for_assign_weight(self.proxy_model, self): - if len(list(model.parameters(recursively=False))) == 0: + if ( + model.model not in self.unassigned_weights_list_recursively + and len(list(model.parameters(recursively=False))) == 0 + ): continue yield model - def traversal_for_auto_layer_map(self): + def traversal_for_layer_map(self): yield self.proxy_model for model in traversal_for_assign_weight(self.proxy_model, self): yield model diff --git a/padiff/report/hooks.py b/padiff/report/hooks.py index d8d2367..2746005 100644 --- a/padiff/report/hooks.py +++ b/padiff/report/hooks.py @@ -22,6 +22,7 @@ for_each_grad_tensor, extract_frame_summary, ) +from paddle.utils import to_sequence import json import numpy import paddle @@ -133,8 +134,7 @@ def info_hook(model, input, output, net_id): # two report_item with same id, the step_idx should be corresponded step_idx = len(list(filter(lambda x: x.type == "forward" and x.net_id == net_id, report.items))) - 1 base_report_node = find_base_report_node(net_id, step_idx) - - retval = map_structure_and_replace_key(replace_forward_output(base_report_node), output, output) + retval = map_structure_and_replace_key(replace_forward_output(base_report_node), to_sequence(output), output) __in_info_hook__ = False return retval else: @@ -158,7 +158,7 @@ def tensor_hook(x_grad, bwd_item, nth_tensor, net_id): ) base_report_node = find_base_report_node(net_id, step_idx) - value = numpy.load(base_report_node["bwd_grads"][nth_tensor]) + value = numpy.load(base_report_node["bwd_grads"][nth_tensor]["path"]) if isinstance(x_grad, paddle.Tensor): return paddle.to_tensor(value) else: @@ -259,7 +259,7 @@ def inner(input_): raise RuntimeError( "In single step mode, try to replace tensor by dumpped numpy value, but the number of tensors and numpy is not equal. Maybe the models are not corresponded." ) - value = numpy.load(numpy_file_list[cur_idx]) + value = numpy.load(numpy_file_list[cur_idx]["path"]) if isinstance(input_, paddle.Tensor): return paddle.to_tensor(value) else: diff --git a/padiff/utils.py b/padiff/utils.py index eb2d77f..d3960f0 100644 --- a/padiff/utils.py +++ b/padiff/utils.py @@ -21,7 +21,7 @@ import torch -from paddle.utils import flatten, pack_sequence_as, map_structure +from paddle.utils import flatten, map_structure, pack_sequence_as """ diff --git a/tests/padiff_unittests.py b/tests/padiff_unittests.py index fec933b..9f88d8c 100644 --- a/tests/padiff_unittests.py +++ b/tests/padiff_unittests.py @@ -13,6 +13,7 @@ # limitations under the License. import os +import sys import subprocess @@ -25,6 +26,7 @@ err_info = f"*** ===================== {fpath} ========================= ***\n" err_info += f"{output}\n" print(f"Failed on unittest {fname} with error message \n {err_info}.", end="\n", flush=True) + sys.exit(1) else: print(f"Succeed on unittest {fname}.", end="\n", flush=True) os.system("rm -rf ./tests/padiff_dump ./tests/padiff_log") diff --git a/tests/test_api_to_Layer.py b/tests/test_api_to_Layer.py index 6a546c6..2e35735 100644 --- a/tests/test_api_to_Layer.py +++ b/tests/test_api_to_Layer.py @@ -61,9 +61,7 @@ def test_api_to_Layer(self): inp = paddle.rand((100, 100), dtype="float32") layer(inp) - layer.report - - assert len(layer.report.items) == 12 + assert len(layer.report.items) == 10 if __name__ == "__main__":