Voice et bot modif

venv/lib/python3.12/site-packages/onnxruntime/quantization/neural_compressor/onnx_model.py

@@ -21,12 +21,12 @@
 import copy
 import logging
 import os
-import sys
 from collections import deque
 from pathlib import Path
 
 import onnx
 import onnx.external_data_helper
+import onnx_ir as ir
 
 from .util import MAXIMUM_PROTOBUF, find_by_name
 
@@ -73,26 +73,11 @@ class ONNXModel:
 
     def check_is_large_model(self):
         """Check model > 2GB."""
-        init_size = 0
-        for init in self._model.graph.initializer:
-            # if initializer has external data location, return True
-            if init.HasField("data_location") and init.data_location == onnx.TensorProto.EXTERNAL:
-                self._is_large_model = True
-                return
-            # if raise error of initializer size > 2GB, return True
-            try:
-                init_bytes = init.SerializeToString()
-                init_size += sys.getsizeof(init_bytes)
-            except Exception as e:
-                if "exceeds maximum protobuf size of 2GB" in str(e):
-                    self._is_large_model = True
-                    return
-                else:  # pragma: no cover
-                    raise e
-            if init_size > MAXIMUM_PROTOBUF:
-                self._is_large_model = True
-                return
-        self._is_large_model = False
+        ir_graph = ir.from_proto(self._model.graph)
+        initializer_size = sum(
+            v.const_value.nbytes for v in ir_graph.initializers.values() if v.const_value is not None
+        )
+        self._is_large_model = initializer_size > MAXIMUM_PROTOBUF
 
     @property
     def is_large_model(self):