File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff line change 1515
1616"""Utility functions related to onnx."""
1717
18+ import copy
1819import io
1920import os
2021import tempfile
@@ -555,14 +556,13 @@ def _get_unique_name(old_name):
555556def check_model (model : onnx .ModelProto ) -> None :
556557 """Checks if the given model is valid."""
557558 if model .ByteSize () > (2 * (1024 ** 3 )): # 2GB limit
558- logger .warning ("Model exceeds 2GB limit, skipping check_model" )
559- # with tempfile.TemporaryDirectory() as temp_dir:
560- # # ONNX also looks in CWD, so we need to use a unique id
561- # unique_id = str(uuid.uuid4())[:8]
562- # onnx_tmp_path = os.path.join(temp_dir, f"model_{unique_id}.onnx")
563- # save_onnx(model, onnx_tmp_path, save_as_external_data=True)
564- # onnx.checker.check_model(onnx_tmp_path)
565-
559+ with tempfile .TemporaryDirectory () as temp_dir :
560+ # ONNX also looks in CWD, so we need to use a unique id
561+ unique_id = str (uuid .uuid4 ())[:8 ]
562+ onnx_tmp_path = os .path .join (temp_dir , f"model_{ unique_id } .onnx" )
563+ model_copy = copy .deepcopy (model )
564+ save_onnx (model_copy , onnx_tmp_path , save_as_external_data = True )
565+ onnx .checker .check_model (onnx_tmp_path )
566566 else :
567567 onnx .checker .check_model (model )
568568
You can’t perform that action at this time.
0 commit comments