已合并
test: Add validation cases for torch.onnx.ONNXProgram.optimize on NPU #43588
木路折创建于 8月2日
test: Add validation cases for torch.onnx.ONNXProgram.optimize on NPU #43588
已合并
共 1 个文件变更+75-0
| @@ -0,0 +1,75 @@ | |||
| 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 direct API validations for ONNXProgram.optimize, so this file is added. | ||
| 19 | +2. This file validates torch.onnx.ONNXProgram.optimize (extendable). | ||
| 20 | +""" | ||
| 21 | + | ||
| 22 | +from __future__ import annotations | ||
| 23 | + | ||
| 24 | +import torch | ||
| 25 | +from torch.testing._internal import common_utils | ||
| 26 | + | ||
| 27 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 28 | + | ||
| 29 | + | ||
| 30 | +class ONNXProgramOptimizeTest(common_utils.TestCase): | ||
| 31 | + """Tests for ONNXProgram.optimize method.""" | ||
| 32 | + | ||
| 33 | + def _create_onnx_program(self) -> torch.onnx.ONNXProgram: | ||
| 34 | + """Helper to create an unoptimized ONNXProgram from a simple model.""" | ||
| 35 | + class Model(torch.nn.Module): | ||
| 36 | + def forward(self, x): | ||
| 37 | + return (x + 1) * 2 | ||
| 38 | + | ||
| 39 | + x = torch.randn(3).to(device_type) | ||
| 40 | + onnx_program = torch.onnx.export( | ||
| 41 | + Model().eval(), (x,), dynamo=True, optimize=False, verbose=False | ||
| 42 | + ) | ||
| 43 | + return onnx_program | ||
| 44 | + | ||
| 45 | + def test_optimize_returns_none(self): | ||
| 46 | + """optimize() should return None.""" | ||
| 47 | + onnx_program = self._create_onnx_program() | ||
| 48 | + result = onnx_program.optimize() | ||
| 49 | + self.assertIsNone(result) | ||
| 50 | + | ||
| 51 | + def test_optimize_model_valid(self): | ||
| 52 | + """The model should be valid after optimize().""" | ||
| 53 | + onnx_program = self._create_onnx_program() | ||
| 54 | + onnx_program.optimize() | ||
| 55 | + self.assertIsNotNone(onnx_program.model) | ||
| 56 | + self.assertIsNotNone(onnx_program.model.graph) | ||
| 57 | + | ||
| 58 | + def test_optimize_idempotent(self): | ||
| 59 | + """Calling optimize() twice should not raise.""" | ||
| 60 | + onnx_program = self._create_onnx_program() | ||
| 61 | + onnx_program.optimize() | ||
| 62 | + onnx_program.optimize() | ||
| 63 | + | ||
| 64 | + def test_optimize_preserves_io_count(self): | ||
| 65 | + """The optimized model should preserve graph input/output count.""" | ||
| 66 | + onnx_program = self._create_onnx_program() | ||
| 67 | + num_inputs = len(onnx_program.model.graph.inputs) | ||
| 68 | + num_outputs = len(onnx_program.model.graph.outputs) | ||
| 69 | + onnx_program.optimize() | ||
| 70 | + self.assertEqual(len(onnx_program.model.graph.inputs), num_inputs) | ||
| 71 | + self.assertEqual(len(onnx_program.model.graph.outputs), num_outputs) | ||
| 72 | + | ||
| 73 | + | ||
| 74 | +if __name__ == "__main__": | ||
| 75 | + common_utils.run_tests() | ||