count size

Author: justinchuby

olive/passes/onnx/common.py

@@ -21,6 +21,8 @@
 
 logger = logging.getLogger(__name__)
 
+_LARGE_IR_MODEL_THRESHOLD = 1536 * 1024 * 1024  # 1536MB
+
 
 def get_external_data_config() -> Dict[str, PassConfigParam]:
     return {
@@ -201,6 +203,11 @@ def model_proto_to_olive_model(
     return olive_model
 
 
+def _count_initializer_size(graph: ir.Graph) -> int:
+    """Count the total size of the initializers in bytes."""
+    return sum(v.const_value.nbytes for v in graph.initializers.values() if v.const_value is not None)
+
+
 def ir_model_to_olive_model(
     model: ir.Model,
     output_model_path: Union[str, Path],
@@ -224,20 +231,28 @@ def ir_model_to_olive_model(
         external_data_config = external_data_config.dict()
 
     save_as_external_data = external_data_config.get("save_as_external_data")
+    # Save as external data if requested or if the model is large
+    # Since we do not have a true estimate of the model architecture size for IR Model,
+    # we count the size of all initializers and limit that to 1.5GB.
+    initializer_size = _count_initializer_size(model.graph)
+    is_large_model = initializer_size > _LARGE_IR_MODEL_THRESHOLD
+    if is_large_model:
+        logger.debug("Model is large (%s), saving as external data", initializer_size)
+    save_as_external_data = save_as_external_data or is_large_model
 
     if save_as_external_data:
         external_data_name = _get_external_data_name(
             Path(output_model_path), external_data_config.get("external_data_name")
         )
         ir.save(model, output_model_path, external_data=external_data_name)
-    else:
-        ir.save(model, output_model_path)
 
-    if external_data_name:
         logger.debug("Model was saved with external data: %s", external_data_name)
         model_path = LocalFolder({"path": Path(output_model_path).parent})
         onnx_file_name = Path(output_model_path).name
+
     else:
+        ir.save(model, output_model_path)
+
         logger.debug("Model was not saved with external data")
         model_path = LocalFile({"path": output_model_path})
         onnx_file_name = None