in optimum/onnx/graph_transformations.py [0:0]
def check_and_save_model(model: onnx.ModelProto, save_path: Optional[Union[str, Path]]):
# We can check ModelProtos that are smaller than 2GB before saving them.
# For larger models, we need to save them first and then check their save path.
# https://github.com/onnx/onnx/blob/main/docs/PythonAPIOverview.md#checking-a-large-onnx-model-2gb
if model.ByteSize() < onnx.checker.MAXIMUM_PROTOBUF:
# For the try catch, refer to https://github.com/microsoft/onnxruntime/issues/14768
try:
onnx.checker.check_model(model)
except Exception as e:
if "No Op registered for" in str(e):
pass
else:
raise e
save_path = Path(save_path).as_posix()
external_file_name = os.path.basename(save_path) + "_data"
external_file_path = os.path.join(os.path.dirname(save_path), external_file_name)
if save_path.endswith(".onnx") and os.path.isfile(save_path):
os.remove(save_path)
model_uses_external_data = False
if os.path.isfile(external_file_path):
model_uses_external_data = True
os.remove(external_file_path)
FORCE_ONNX_EXTERNAL_DATA = os.getenv("FORCE_ONNX_EXTERNAL_DATA", "0") == "1"
onnx.save(
model,
save_path,
save_as_external_data=model_uses_external_data or FORCE_ONNX_EXTERNAL_DATA,
all_tensors_to_one_file=True,
location=external_file_name,
convert_attribute=True,
size_threshold=1024 if not FORCE_ONNX_EXTERNAL_DATA else 100,
)
try:
onnx.checker.check_model(save_path)
except Exception as e:
if "No Op registered for" in str(e):
pass
else:
raise e