diff --git a/onnxruntime/python/tools/quantization/quant_utils.py b/onnxruntime/python/tools/quantization/quant_utils.py index eeef048378..2ceefeadcd 100644 --- a/onnxruntime/python/tools/quantization/quant_utils.py +++ b/onnxruntime/python/tools/quantization/quant_utils.py @@ -545,10 +545,10 @@ def model_has_infer_metadata(model): def load_model_with_shape_infer(model_path: Path): - with tempfile.TemporaryDirectory() as temp_dir: - inferred_model_path = str(Path(temp_dir) / (model_path.stem + "-inferred.onnx")) - onnx.shape_inference.infer_shapes_path(str(model_path), inferred_model_path) - model = onnx.load(inferred_model_path) + inferred_model_path = generate_identified_filename(model_path, "-inferred") + onnx.shape_inference.infer_shapes_path(str(model_path), str(inferred_model_path)) + model = onnx.load(inferred_model_path.as_posix()) + inferred_model_path.unlink() return model diff --git a/onnxruntime/test/python/quantization/test_quant_util.py b/onnxruntime/test/python/quantization/test_quant_util.py index d7da0b8923..5a9a5a5029 100644 --- a/onnxruntime/test/python/quantization/test_quant_util.py +++ b/onnxruntime/test/python/quantization/test_quant_util.py @@ -6,11 +6,15 @@ # license information. # -------------------------------------------------------------------------- +import tempfile import unittest +from pathlib import Path import numpy +import onnx +from onnx import TensorProto, helper, numpy_helper -from onnxruntime.quantization.quant_utils import compute_scale_zp +from onnxruntime.quantization.quant_utils import compute_scale_zp, load_model, model_has_infer_metadata class TestQuantUtil(unittest.TestCase): @@ -30,6 +34,37 @@ class TestQuantUtil(unittest.TestCase): self.assertEqual(compute_scale_zp(-tiny_float, tiny_float, 0, 255, symmetric=True), [0, 1.0]) self.assertEqual(compute_scale_zp(-tiny_float, 0.0, 0, 255, symmetric=False), [0, 1.0]) + def test_load_external_model(self): + input_name = "input" + output_name = "output" + add_shape = [1024, 1024] + + initializers = [] + weight_name = "weight" + weight_data = numpy.random.normal(0, 0.1, add_shape).astype(numpy.float32) + initializers.append(numpy_helper.from_array(weight_data, name=weight_name)) + add_node = helper.make_node("Add", [input_name, weight_name], [output_name], name="add_node") + + # make graph + input_tensor = helper.make_tensor_value_info(input_name, TensorProto.FLOAT, add_shape) + output_tensor = helper.make_tensor_value_info(output_name, TensorProto.FLOAT, add_shape) + graph_name = "test_load_external_model" + graph = helper.make_graph( + [add_node], + graph_name, + [input_tensor], + [output_tensor], + initializer=initializers, + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + + with tempfile.TemporaryDirectory() as temp_dir: + self.assertFalse(model_has_infer_metadata(model)) + model_file_path = temp_dir + "/test_load_external_model.onnx" + onnx.save(model, model_file_path, save_as_external_data=True) + model_reloaded = load_model(Path(model_file_path), False) + self.assertTrue(model_has_infer_metadata(model_reloaded)) + if __name__ == "__main__": unittest.main()