-
Notifications
You must be signed in to change notification settings - Fork 57
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
67d8ab9
commit 357c41b
Showing
1 changed file
with
98 additions
and
106 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,128 +1,120 @@ | ||
| # Copyright (c) Microsoft Corporation. | ||
| # Licensed under the MIT License. | ||
| """Unit tests for the _io module.""" | ||
|
|
||
| import os | ||
| import tempfile | ||
| import unittest | ||
|
|
||
| import numpy as np | ||
|
|
||
| from onnxscript import ir | ||
| from onnxscript.ir import _io | ||
|
|
||
|
|
||
| def _create_simple_model(): | ||
| tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X") | ||
| node = ir.Node("Identity", inputs=[tensor], outputs=["Y"]) | ||
| graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor]) | ||
| return ir.model(graph) | ||
|
|
||
| class TestIOFunctions(unittest.TestCase): | ||
| def _create_initializer(tensor: ir.TensorProtocol) -> ir.Value: | ||
| return ir.Value( | ||
| name=tensor.name, | ||
| shape=tensor.shape, | ||
| type=ir.TensorType(tensor.dtype), | ||
| const_value=tensor, | ||
| ) | ||
|
|
||
|
|
||
| def _create_simple_model_with_initializers() -> ir.Model: | ||
| tensor_0 = ir.tensor([0.0], dtype=ir.DataType.FLOAT, name="initializer_0") | ||
| initializer = _create_initializer(tensor_0) | ||
| tensor_1 = ir.tensor([1.0], dtype=ir.DataType.FLOAT) | ||
| identity_node = ir.Node("", "Identity", inputs=(initializer,)) | ||
| identity_node.outputs[0].shape = ir.Shape([1]) | ||
| identity_node.outputs[0].dtype = ir.DataType.FLOAT | ||
| identity_node.outputs[0].name = "identity_0" | ||
| const_node = ir.Node( | ||
| "", | ||
| "Constant", | ||
| inputs=(), | ||
| outputs=( | ||
| ir.Value(name="const_0", shape=tensor_1.shape, type=ir.TensorType(tensor_1.dtype)), | ||
| ), | ||
| attributes=ir.convenience.convert_attributes(dict(value=tensor_1)), | ||
| ) | ||
| graph = ir.Graph( | ||
| inputs=[initializer], | ||
| outputs=[*identity_node.outputs, *const_node.outputs], | ||
| nodes=[identity_node, const_node], | ||
| initializers=[initializer], | ||
| name="test_graph", | ||
| ) | ||
| print(graph) | ||
| return ir.Model(graph, ir_version=10) | ||
|
|
||
|
|
||
| class IOFunctionsTest(unittest.TestCase): | ||
| def test_load(self): | ||
| model = _create_simple_model_with_initializers() | ||
| with tempfile.TemporaryDirectory() as tmpdir: | ||
| path = os.path.join(tmpdir, "model.onnx") | ||
| # Create a simple ONNX model | ||
| tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X") | ||
| node = ir.node("Identity", inputs=[tensor], outputs=["Y"]) | ||
| graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor]) | ||
| model = ir.model(graph) | ||
| # Save the model to a file | ||
| with open(path, "wb") as f: | ||
| f.write(model.SerializeToString()) | ||
|
|
||
| # Load the model using the _io.load function | ||
| _io.save(model, path) | ||
| loaded_model = _io.load(path) | ||
|
|
||
| # Check that the loaded model is correct | ||
| self.assertEqual(loaded_model.graph.name, "test_graph") | ||
| self.assertEqual(len(loaded_model.graph.initializers), 1) | ||
| self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0]) | ||
|
|
||
| def test_save(self): | ||
| self.assertEqual(loaded_model.ir_version, model.ir_version) | ||
| self.assertEqual(loaded_model.graph.name, model.graph.name) | ||
| self.assertEqual(len(loaded_model.graph.initializers), 1) | ||
| self.assertEqual(len(loaded_model.graph), 2) | ||
| np.testing.assert_array_equal( | ||
| loaded_model.graph.initializers["initializer_0"].const_value.numpy(), | ||
| np.array([0.0]), | ||
| ) | ||
| np.testing.assert_array_equal( | ||
| loaded_model.graph.node(1).attributes["value"].as_tensor().numpy(), np.array([1.0]) | ||
| ) | ||
| self.assertEqual(loaded_model.graph.inputs[0].name, "initializer_0") | ||
| self.assertEqual(loaded_model.graph.outputs[0].name, "identity_0") | ||
| self.assertEqual(loaded_model.graph.outputs[1].name, "const_0") | ||
|
|
||
| def test_save_with_external_data_modify_model_false_does_not_modify_model(self): | ||
| model = _create_simple_model_with_initializers() | ||
| self.assertIsInstance(model.graph.initializers["initializer_0"].const_value, ir.Tensor) | ||
| with tempfile.TemporaryDirectory() as tmpdir: | ||
| path = os.path.join(tmpdir, "model.onnx") | ||
| external_data_path = "external_data" | ||
|
|
||
| # Create a simple ONNX model | ||
| tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X") | ||
| node = ir.node("Identity", inputs=[tensor], outputs=["Y"]) | ||
| graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor]) | ||
| model = ir.model(graph) | ||
| core_model = _core.Model(model) | ||
|
|
||
| # Save the model using the _io.save function | ||
| _io.save(core_model, path, external_data=external_data_path, modify_model=True) | ||
|
|
||
| # Load the model back to verify it was saved correctly | ||
| external_data_path = "model.data" | ||
| _io.save(model, path, external_data=external_data_path) | ||
| loaded_model = _io.load(path) | ||
|
|
||
| # Check that the loaded model is correct | ||
| self.assertEqual(loaded_model.graph.name, "test_graph") | ||
| self.assertEqual(len(loaded_model.graph.initializers), 1) | ||
| self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0]) | ||
|
|
||
| def test_save_without_external_data(self): | ||
| # The loaded model contains external data | ||
| initializer_tensor = loaded_model.graph.initializers["initializer_0"].const_value | ||
| self.assertIsInstance(initializer_tensor, ir.ExternalTensor) | ||
| # The attribute is not externalized | ||
| const_attr_tensor = loaded_model.graph.node(1).attributes["value"].as_tensor() | ||
| self.assertIsInstance(const_attr_tensor, ir.TensorProtoTensor) | ||
| np.testing.assert_array_equal(initializer_tensor.numpy(), np.array([0.0])) | ||
| np.testing.assert_array_equal(const_attr_tensor.numpy(), np.array([1.0])) | ||
|
|
||
| # The original model is not changed and can be accessed even if the | ||
| # external data file is deleted | ||
| initializer_tensor = model.graph.initializers["initializer_0"].const_value | ||
| self.assertIsInstance(initializer_tensor, ir.Tensor) | ||
| const_attr_tensor = model.graph.node(1).attributes["value"].as_tensor() | ||
| self.assertIsInstance(const_attr_tensor, ir.Tensor) | ||
| np.testing.assert_array_equal(initializer_tensor.numpy(), np.array([0.0])) | ||
| np.testing.assert_array_equal(const_attr_tensor.numpy(), np.array([1.0])) | ||
|
|
||
| def test_save_with_external_data_modify_model(self): | ||
| model = _create_simple_model_with_initializers() | ||
| self.assertIsInstance(model.graph.initializers["initializer_0"].const_value, ir.Tensor) | ||
| with tempfile.TemporaryDirectory() as tmpdir: | ||
| path = os.path.join(tmpdir, "model.onnx") | ||
| external_data_path = "model.data" | ||
| _io.save(model, path, external_data=external_data_path, modify_model=True) | ||
|
|
||
| # The original model is modified | ||
| initializer_tensor = model.graph.initializers["initializer_0"].const_value | ||
| self.assertIsInstance(initializer_tensor, ir.ExternalTensor) | ||
| const_attr_tensor = model.graph.node(1).attributes["value"].as_tensor() | ||
| # But the attribute is not externalized | ||
| self.assertIsInstance(const_attr_tensor, ir.Tensor) | ||
| np.testing.assert_array_equal(initializer_tensor.numpy(), np.array([0.0])) | ||
| np.testing.assert_array_equal(const_attr_tensor.numpy(), np.array([1.0])) | ||
|
|
||
| # Create a simple ONNX model | ||
| tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X") | ||
| node = ir.node("Identity", inputs=[tensor], outputs=["Y"]) | ||
| graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor]) | ||
| model = ir.model(graph) | ||
| core_model = _core.Model(model) | ||
|
|
||
| # Save the model using the _io.save function without external data | ||
| _io.save(core_model, path, modify_model=True) | ||
|
|
||
| # Load the model back to verify it was saved correctly | ||
| loaded_model = _io.load(path) | ||
|
|
||
| # Check that the loaded model is correct | ||
| self.assertEqual(loaded_model.graph.name, "test_graph") | ||
| self.assertEqual(len(loaded_model.graph.initializers), 1) | ||
| self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0]) | ||
|
|
||
| def test_save_with_external_data_modify_model_true(self): | ||
| with tempfile.TemporaryDirectory() as tmpdir: | ||
| path = os.path.join(tmpdir, "model.onnx") | ||
| external_data_path = "external_data" | ||
|
|
||
| # Create a simple ONNX model | ||
| tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X") | ||
| node = ir.node("Identity", inputs=[tensor], outputs=["Y"]) | ||
| graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor]) | ||
| model = ir.model(graph) | ||
| core_model = _core.Model(model) | ||
|
|
||
| # Save the model using the _io.save function with external data and modify_model=True | ||
| _io.save(core_model, path, external_data=external_data_path, modify_model=True) | ||
|
|
||
| # Load the model back to verify it was saved correctly | ||
| loaded_model = _io.load(path) | ||
|
|
||
| # Check that the loaded model is correct | ||
| self.assertEqual(loaded_model.graph.name, "test_graph") | ||
| self.assertEqual(len(loaded_model.graph.initializers), 1) | ||
| self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0]) | ||
|
|
||
| def test_save_with_external_data_modify_model_false(self): | ||
| with tempfile.TemporaryDirectory() as tmpdir: | ||
| path = os.path.join(tmpdir, "model.onnx") | ||
| external_data_path = "external_data" | ||
|
|
||
| # Create a simple ONNX model | ||
| tensor = ir.tensor([1.0], dtype=ir.DataType.FLOAT, name="X") | ||
| node = ir.node("Identity", inputs=[tensor], outputs=["Y"]) | ||
| graph = ir.graph([node], name="test_graph", outputs=[node.outputs[0]], initializers=[tensor]) | ||
| model = ir.model(graph) | ||
| core_model = _core.Model(model) | ||
|
|
||
| # Save the model using the _io.save function with external data and modify_model=False | ||
| _io.save(core_model, path, external_data=external_data_path, modify_model=False) | ||
|
|
||
| # Load the model back to verify it was saved correctly | ||
| loaded_model = _io.load(path) | ||
|
|
||
| # Check that the loaded model is correct | ||
| self.assertEqual(loaded_model.graph.name, "test_graph") | ||
| self.assertEqual(len(loaded_model.graph.initializers), 1) | ||
| self.assertEqual(loaded_model.graph.initializers["X"].const_value, [1.0]) | ||
|
|
||
| if __name__ == "__main__": | ||
| unittest.main() |