已合并
test: add test cases for torch.onnx.ONNXProgram.model_proto #43581
木路折创建于 8月2日
test: add test cases for torch.onnx.ONNXProgram.model_proto #43581
已合并
共 1 个文件变更+87-0
| @@ -0,0 +1,87 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd | ||
| 2 | +# All rights reserved. | ||
| 3 | +# | ||
| 4 | +# Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +# you may not use this file except in compliance with the License. | ||
| 6 | +# You may obtain a copy of the License at | ||
| 7 | +# | ||
| 8 | +# https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +# | ||
| 10 | +# Unless required by applicable law or agreed to in writing, software | ||
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +# See the License for the specific language governing permissions and | ||
| 14 | +# limitations under the License. | ||
| 15 | + | ||
| 16 | +""" | ||
| 17 | +Add validation cases for torch.onnx.ONNXProgram APIs on NPU: | ||
| 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.model_proto (extendable). | ||
| 20 | +""" | ||
| 21 | + | ||
| 22 | +import torch | ||
| 23 | +from torch.testing._internal import common_utils | ||
| 24 | + | ||
| 25 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 26 | + | ||
| 27 | + | ||
| 28 | +class TestONNXProgramModelProto(common_utils.TestCase): | ||
| 29 | + | ||
| 30 | + def test_model_proto_returns_valid_proto(self): | ||
| 31 | + class Model(torch.nn.Module): | ||
| 32 | + def forward(self, x): | ||
| 33 | + return x + 1 | ||
| 34 | + | ||
| 35 | + x = torch.randn(3).to(device_type) | ||
| 36 | + onnx_program = torch.onnx.export(Model(), (x,), dynamo=True) | ||
| 37 | + proto = onnx_program.model_proto | ||
| 38 | + self.assertIsNotNone(proto) | ||
| 39 | + self.assertGreater(proto.ir_version, 0) | ||
| 40 | + | ||
| 41 | + def test_model_proto_graph_structure(self): | ||
| 42 | + class Model(torch.nn.Module): | ||
| 43 | + def forward(self, x): | ||
| 44 | + return x * 2 | ||
| 45 | + | ||
| 46 | + x = torch.randn(3).to(device_type) | ||
| 47 | + onnx_program = torch.onnx.export(Model(), (x,), dynamo=True) | ||
| 48 | + proto = onnx_program.model_proto | ||
| 49 | + self.assertEqual(len(proto.graph.input), 1) | ||
| 50 | + self.assertEqual(len(proto.graph.output), 1) | ||
| 51 | + | ||
| 52 | + def test_model_proto_producer_name(self): | ||
| 53 | + class Model(torch.nn.Module): | ||
| 54 | + def forward(self, x): | ||
| 55 | + return x | ||
| 56 | + | ||
| 57 | + x = torch.randn(3).to(device_type) | ||
| 58 | + onnx_program = torch.onnx.export(Model(), (x,), dynamo=True) | ||
| 59 | + proto = onnx_program.model_proto | ||
| 60 | + self.assertEqual(proto.producer_name, "pytorch") | ||
| 61 | + | ||
| 62 | + def test_model_proto_serialization(self): | ||
| 63 | + class Model(torch.nn.Module): | ||
| 64 | + def forward(self, x): | ||
| 65 | + return x + x | ||
| 66 | + | ||
| 67 | + x = torch.randn(3).to(device_type) | ||
| 68 | + onnx_program = torch.onnx.export(Model(), (x,), dynamo=True) | ||
| 69 | + proto = onnx_program.model_proto | ||
| 70 | + serialized = proto.SerializeToString() | ||
| 71 | + self.assertGreater(len(serialized), 0) | ||
| 72 | + | ||
| 73 | + def test_model_proto_multiple_inputs_outputs(self): | ||
| 74 | + class Model(torch.nn.Module): | ||
| 75 | + def forward(self, x, y): | ||
| 76 | + return x + y, x - y | ||
| 77 | + | ||
| 78 | + x = torch.randn(3).to(device_type) | ||
| 79 | + y = torch.randn(3).to(device_type) | ||
| 80 | + onnx_program = torch.onnx.export(Model(), (x, y), dynamo=True) | ||
| 81 | + proto = onnx_program.model_proto | ||
| 82 | + self.assertEqual(len(proto.graph.input), 2) | ||
| 83 | + self.assertEqual(len(proto.graph.output), 2) | ||
| 84 | + | ||
| 85 | + | ||
| 86 | +if __name__ == "__main__": | ||
| 87 | + common_utils.run_tests() | ||