Skip to content

Commit f7c3f5a

Browse files
justinchubyCopilot
andauthored
[IR] Handle external initializers in subgraphs (#2347)
Support converting external initializers in subgraphs. --------- Co-authored-by: Copilot <[email protected]>
1 parent 024a9cd commit f7c3f5a

File tree

1 file changed

+21
-15
lines changed

1 file changed

+21
-15
lines changed

onnxscript/ir/external_data.py

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -306,17 +306,20 @@ def convert_tensors_to_external(
306306
def load_to_model(model: _core.Model) -> _core.Model:
307307
"""Convert all external model initializers to memory tensors in-place.
308308
309+
All initializers in the main graph and subgraphs are handled.
310+
309311
Args:
310312
model: Model to process.
311313
"""
312-
# TODO(justinchuby): Load attributes and initializers in subgraphs
314+
# TODO(justinchuby): Load tensor attributes in subgraphs
313315
values_to_convert = []
314-
for value in model.graph.initializers.values():
315-
if value.const_value is None:
316-
# Filter out the uninitialized initializer values
317-
continue
318-
if isinstance(value.const_value, _core.ExternalTensor):
319-
values_to_convert.append(value)
316+
for graph in model.graphs():
317+
for value in graph.initializers.values():
318+
if value.const_value is None:
319+
# Filter out the uninitialized initializer values
320+
continue
321+
if isinstance(value.const_value, _core.ExternalTensor):
322+
values_to_convert.append(value)
320323
loaded_tensors = convert_tensors_from_external(
321324
[v.const_value for v in values_to_convert] # type: ignore[misc]
322325
)
@@ -346,6 +349,8 @@ def unload_from_model(
346349
to load the newly saved model, or provide a different external data path that
347350
is not currently referenced by any tensors in the model.
348351
352+
All initializers in the main graph and subgraphs are handled.
353+
349354
Args:
350355
model: Model to process.
351356
base_dir: Path the directory where the ONNX model file is.
@@ -361,14 +366,15 @@ def unload_from_model(
361366
initializers_to_become_external = []
362367
# Existing external tensors, if below the threshold, should be loaded to memory
363368
initializers_to_load_to_memory = []
364-
for value in model.graph.initializers.values():
365-
if value.const_value is None:
366-
# Filter out the uninitialized initializer values
367-
continue
368-
if value.const_value.nbytes > size_threshold_bytes:
369-
initializers_to_become_external.append(value)
370-
elif isinstance(value.const_value, _core.ExternalTensor):
371-
initializers_to_load_to_memory.append(value)
369+
for graph in model.graphs():
370+
for value in graph.initializers.values():
371+
if value.const_value is None:
372+
# Filter out the uninitialized initializer values
373+
continue
374+
if value.const_value.nbytes > size_threshold_bytes:
375+
initializers_to_become_external.append(value)
376+
elif isinstance(value.const_value, _core.ExternalTensor):
377+
initializers_to_load_to_memory.append(value)
372378

373379
# Load to memory first, then convert to external tensors, because
374380
# the existing external tensors may be overwritten by the new external data

0 commit comments

Comments
 (0)