From 6a02b10b474627cc7ad0f77a2db5161d3f68c0e7 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 May 2025 12:30:50 -0700 Subject: [PATCH 1/6] [IR] Handle initializers in subgraphs --- onnxscript/ir/external_data.py | 48 ++++++++++++++++++++--------- onnxscript/ir/external_data_test.py | 1 - 2 files changed, 33 insertions(+), 16 deletions(-) diff --git a/onnxscript/ir/external_data.py b/onnxscript/ir/external_data.py index 4cf6d72f91..bb34eeb526 100644 --- a/onnxscript/ir/external_data.py +++ b/onnxscript/ir/external_data.py @@ -306,17 +306,20 @@ def convert_tensors_to_external( def load_to_model(model: _core.Model) -> _core.Model: """Convert all external model initializers to memory tensors in-place. + All initializers in the main graph and subgraphs are be handled. + Args: model: Model to process. """ - # TODO(justinchuby): Load attributes and initializers in subgraphs + # TODO(justinchuby): Load tensor attributes in subgraphs values_to_convert = [] - for value in model.graph.initializers.values(): - if value.const_value is None: - # Filter out the uninitialized initializer values - continue - if isinstance(value.const_value, _core.ExternalTensor): - values_to_convert.append(value) + for graph in model.graphs(): + for value in graph.initializers.values(): + if value.const_value is None: + # Filter out the uninitialized initializer values + continue + if isinstance(value.const_value, _core.ExternalTensor): + values_to_convert.append(value) loaded_tensors = convert_tensors_from_external( [v.const_value for v in values_to_convert] # type: ignore[misc] ) @@ -346,6 +349,8 @@ def unload_from_model( to load the newly saved model, or provide a different external data path that is not currently referenced by any tensors in the model. + All initializers in the main graph and subgraphs are be handled. + Args: model: Model to process. base_dir: Path the directory where the ONNX model file is. @@ -361,14 +366,27 @@ def unload_from_model( initializers_to_become_external = [] # Existing external tensors, if below the threshold, should be loaded to memory initializers_to_load_to_memory = [] - for value in model.graph.initializers.values(): - if value.const_value is None: - # Filter out the uninitialized initializer values - continue - if value.const_value.nbytes > size_threshold_bytes: - initializers_to_become_external.append(value) - elif isinstance(value.const_value, _core.ExternalTensor): - initializers_to_load_to_memory.append(value) + seen_names: set[str] = set() + for graph in model.graphs(): + for value in graph.initializers.values(): + if value.name is None: + raise ValueError( + "Initializer name cannot be None. " + "Please ensure all initializers have unique names." + ) + if value.name in seen_names: + raise ValueError( + f"Initializer name '{value.name}' is not unique in the model. " + "Please ensure all initializers have unique names." + ) + seen_names.add(value.name) + if value.const_value is None: + # Filter out the uninitialized initializer values + continue + if value.const_value.nbytes > size_threshold_bytes: + initializers_to_become_external.append(value) + elif isinstance(value.const_value, _core.ExternalTensor): + initializers_to_load_to_memory.append(value) # Load to memory first, then convert to external tensors, because # the existing external tensors may be overwritten by the new external data diff --git a/onnxscript/ir/external_data_test.py b/onnxscript/ir/external_data_test.py index 11de6285c9..3789a4b5bc 100644 --- a/onnxscript/ir/external_data_test.py +++ b/onnxscript/ir/external_data_test.py @@ -497,6 +497,5 @@ def test_external_data_sorted(self): self.assertEqual(tensor_data, tensor_bytes) self.assertEqual(tensor_data, expected_tensor_order[i]) - if __name__ == "__main__": unittest.main() From e5b403c259d6a5c4635d0fe3f96094c787e8826e Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 May 2025 13:11:20 -0700 Subject: [PATCH 2/6] Update onnxscript/ir/external_data.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- onnxscript/ir/external_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxscript/ir/external_data.py b/onnxscript/ir/external_data.py index bb34eeb526..436064dbe6 100644 --- a/onnxscript/ir/external_data.py +++ b/onnxscript/ir/external_data.py @@ -306,7 +306,7 @@ def convert_tensors_to_external( def load_to_model(model: _core.Model) -> _core.Model: """Convert all external model initializers to memory tensors in-place. - All initializers in the main graph and subgraphs are be handled. + All initializers in the main graph and subgraphs are handled. Args: model: Model to process. From 907eee6177bd82b365821003aa858fc711c5f3aa Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 May 2025 13:11:28 -0700 Subject: [PATCH 3/6] Update onnxscript/ir/external_data.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- onnxscript/ir/external_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxscript/ir/external_data.py b/onnxscript/ir/external_data.py index 436064dbe6..0c23181e19 100644 --- a/onnxscript/ir/external_data.py +++ b/onnxscript/ir/external_data.py @@ -349,7 +349,7 @@ def unload_from_model( to load the newly saved model, or provide a different external data path that is not currently referenced by any tensors in the model. - All initializers in the main graph and subgraphs are be handled. + All initializers in the main graph and subgraphs are handled. Args: model: Model to process. From 5db2ac51cc8b3804046503491a988d537997c96b Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 May 2025 13:58:54 -0700 Subject: [PATCH 4/6] Do not check this --- onnxscript/ir/external_data.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/onnxscript/ir/external_data.py b/onnxscript/ir/external_data.py index 0c23181e19..a864b8f10a 100644 --- a/onnxscript/ir/external_data.py +++ b/onnxscript/ir/external_data.py @@ -366,7 +366,6 @@ def unload_from_model( initializers_to_become_external = [] # Existing external tensors, if below the threshold, should be loaded to memory initializers_to_load_to_memory = [] - seen_names: set[str] = set() for graph in model.graphs(): for value in graph.initializers.values(): if value.name is None: @@ -374,12 +373,6 @@ def unload_from_model( "Initializer name cannot be None. " "Please ensure all initializers have unique names." ) - if value.name in seen_names: - raise ValueError( - f"Initializer name '{value.name}' is not unique in the model. " - "Please ensure all initializers have unique names." - ) - seen_names.add(value.name) if value.const_value is None: # Filter out the uninitialized initializer values continue From f7ec27a403536f1a4bc6ec49d16cdd4a7066df4d Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 May 2025 14:14:09 -0700 Subject: [PATCH 5/6] name --- onnxscript/ir/external_data.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/onnxscript/ir/external_data.py b/onnxscript/ir/external_data.py index a864b8f10a..af0bb226ca 100644 --- a/onnxscript/ir/external_data.py +++ b/onnxscript/ir/external_data.py @@ -368,11 +368,6 @@ def unload_from_model( initializers_to_load_to_memory = [] for graph in model.graphs(): for value in graph.initializers.values(): - if value.name is None: - raise ValueError( - "Initializer name cannot be None. " - "Please ensure all initializers have unique names." - ) if value.const_value is None: # Filter out the uninitialized initializer values continue From adca07953824b43b79a56082d89140406c67438f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 27 May 2025 16:00:12 -0700 Subject: [PATCH 6/6] lint --- onnxscript/ir/external_data_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/onnxscript/ir/external_data_test.py b/onnxscript/ir/external_data_test.py index 3789a4b5bc..11de6285c9 100644 --- a/onnxscript/ir/external_data_test.py +++ b/onnxscript/ir/external_data_test.py @@ -497,5 +497,6 @@ def test_external_data_sorted(self): self.assertEqual(tensor_data, tensor_bytes) self.assertEqual(tensor_data, expected_tensor_order[i]) + if __name__ == "__main__": unittest.main()