File tree Expand file tree Collapse file tree 1 file changed +22
-3
lines changed
tools/onnx-graphsurgeon/onnx_graphsurgeon/ir Expand file tree Collapse file tree 1 file changed +22
-3
lines changed Original file line number Diff line number Diff line change @@ -1245,12 +1245,31 @@ def should_eval_foldable(tensor):
1245
1245
else :
1246
1246
names = [t .name for t in graph_clone .outputs ]
1247
1247
try :
1248
+ import os
1249
+ import tempfile
1250
+ import onnx
1248
1251
import onnxruntime as onnxrt
1249
1252
1253
+ onnx_model = export_onnx (graph_clone , do_type_check = False )
1254
+ if onnx_model .ByteSize () >= onnx .checker .MAXIMUM_PROTOBUF :
1255
+ tmp_dir = tempfile .TemporaryDirectory ()
1256
+ tmp_path = os .path .join (tmp_dir .name , "tmp.onnx" )
1257
+ location = os .path .basename (tmp_path ) + ".data"
1258
+ if os .path .exists (location ):
1259
+ os .remove (location )
1260
+ onnx .save (
1261
+ onnx_model ,
1262
+ tmp_path ,
1263
+ save_as_external_data = True ,
1264
+ all_tensors_to_one_file = True ,
1265
+ location = location ,
1266
+ )
1267
+ onnx_model = tmp_path
1268
+ else :
1269
+ onnx_model = onnx_model .SerializeToString ()
1270
+
1250
1271
sess = onnxrt .InferenceSession (
1251
- export_onnx (
1252
- graph_clone , do_type_check = False
1253
- ).SerializeToString (),
1272
+ onnx_model ,
1254
1273
providers = ORT_PROVIDERS ,
1255
1274
)
1256
1275
values = sess .run (names , {})
You can’t perform that action at this time.
0 commit comments