@@ -306,17 +306,20 @@ def convert_tensors_to_external(
306
306
def load_to_model (model : _core .Model ) -> _core .Model :
307
307
"""Convert all external model initializers to memory tensors in-place.
308
308
309
+ All initializers in the main graph and subgraphs are handled.
310
+
309
311
Args:
310
312
model: Model to process.
311
313
"""
312
- # TODO(justinchuby): Load attributes and initializers in subgraphs
314
+ # TODO(justinchuby): Load tensor attributes in subgraphs
313
315
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 )
320
323
loaded_tensors = convert_tensors_from_external (
321
324
[v .const_value for v in values_to_convert ] # type: ignore[misc]
322
325
)
@@ -346,6 +349,8 @@ def unload_from_model(
346
349
to load the newly saved model, or provide a different external data path that
347
350
is not currently referenced by any tensors in the model.
348
351
352
+ All initializers in the main graph and subgraphs are handled.
353
+
349
354
Args:
350
355
model: Model to process.
351
356
base_dir: Path the directory where the ONNX model file is.
@@ -361,14 +366,15 @@ def unload_from_model(
361
366
initializers_to_become_external = []
362
367
# Existing external tensors, if below the threshold, should be loaded to memory
363
368
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 )
372
378
373
379
# Load to memory first, then convert to external tensors, because
374
380
# the existing external tensors may be overwritten by the new external data
0 commit comments