已开启
test: add test cases for torch.onnx.ONNXProgram.model_proto #43582
木路折创建于 8月2日
test: add test cases for torch.onnx.ONNXProgram.model_proto #43582
已开启
共 1 个文件变更+60-2
| @@ -15,8 +15,8 @@ | |||
| 15 | 15 | ||
| 16 | """ | 16 | """ |
| 17 | Add validation cases for torch.onnx.ONNXProgram APIs on NPU: | 17 | Add validation cases for torch.onnx.ONNXProgram APIs on NPU: |
| 18 | -1. PyTorch community lacks direct API validations for ONNXProgram.optimize, so this file is added. | 18 | +1. PyTorch community lacks sufficient and direct API validations for some APIs, so this file is added. |
| 19 | -2. This file validates torch.onnx.ONNXProgram.optimize (extendable). | 19 | +2. This file validates torch.onnx.ONNXProgram.optimize, torch.onnx.ONNXProgram.model_proto (extendable). |
| 20 | """ | 20 | """ |
| 21 | 21 | ||
| 22 | from __future__ import annotations | 22 | from __future__ import annotations |
| @@ -71,5 +71,63 @@ class ONNXProgramOptimizeTest(common_utils.TestCase): | |||
| 71 | self.assertEqual(len(onnx_program.model.graph.outputs), num_outputs) | 71 | self.assertEqual(len(onnx_program.model.graph.outputs), num_outputs) |
| 72 | 72 | ||
| 73 | 73 | ||
| 74 | +class TestONNXProgramModelProto(common_utils.TestCase): | ||
| 75 | + | ||
| 76 | + def test_model_proto_returns_valid_proto(self): | ||
| 77 | + class Model(torch.nn.Module): | ||
| 78 | + def forward(self, x): | ||
| 79 | + return x + 1 | ||
| 80 | + | ||
| 81 | + x = torch.randn(3).to(device_type) | ||
| 82 | + onnx_program = torch.onnx.export(Model(), (x,), dynamo=True) | ||
| 83 | + proto = onnx_program.model_proto | ||
| 84 | + self.assertIsNotNone(proto) | ||
| 85 | + self.assertGreater(proto.ir_version, 0) | ||
| 86 | + | ||
| 87 | + def test_model_proto_graph_structure(self): | ||
| 88 | + class Model(torch.nn.Module): | ||
| 89 | + def forward(self, x): | ||
| 90 | + return x * 2 | ||
| 91 | + | ||
| 92 | + x = torch.randn(3).to(device_type) | ||
| 93 | + onnx_program = torch.onnx.export(Model(), (x,), dynamo=True) | ||
| 94 | + proto = onnx_program.model_proto | ||
| 95 | + self.assertEqual(len(proto.graph.input), 1) | ||
| 96 | + self.assertEqual(len(proto.graph.output), 1) | ||
| 97 | + | ||
| 98 | + def test_model_proto_producer_name(self): | ||
| 99 | + class Model(torch.nn.Module): | ||
| 100 | + def forward(self, x): | ||
| 101 | + return x | ||
| 102 | + | ||
| 103 | + x = torch.randn(3).to(device_type) | ||
| 104 | + onnx_program = torch.onnx.export(Model(), (x,), dynamo=True) | ||
| 105 | + proto = onnx_program.model_proto | ||
| 106 | + self.assertEqual(proto.producer_name, "pytorch") | ||
| 107 | + | ||
| 108 | + def test_model_proto_serialization(self): | ||
| 109 | + class Model(torch.nn.Module): | ||
| 110 | + def forward(self, x): | ||
| 111 | + return x + x | ||
| 112 | + | ||
| 113 | + x = torch.randn(3).to(device_type) | ||
| 114 | + onnx_program = torch.onnx.export(Model(), (x,), dynamo=True) | ||
| 115 | + proto = onnx_program.model_proto | ||
| 116 | + serialized = proto.SerializeToString() | ||
| 117 | + self.assertGreater(len(serialized), 0) | ||
| 118 | + | ||
| 119 | + def test_model_proto_multiple_inputs_outputs(self): | ||
| 120 | + class Model(torch.nn.Module): | ||
| 121 | + def forward(self, x, y): | ||
| 122 | + return x + y, x - y | ||
| 123 | + | ||
| 124 | + x = torch.randn(3).to(device_type) | ||
| 125 | + y = torch.randn(3).to(device_type) | ||
| 126 | + onnx_program = torch.onnx.export(Model(), (x, y), dynamo=True) | ||
| 127 | + proto = onnx_program.model_proto | ||
| 128 | + self.assertEqual(len(proto.graph.input), 2) | ||
| 129 | + self.assertEqual(len(proto.graph.output), 2) | ||
| 130 | + | ||
| 131 | + | ||
| 74 | if __name__ == "__main__": | 132 | if __name__ == "__main__": |
| 75 | common_utils.run_tests() | 133 | common_utils.run_tests() |