已开启
test: Add validation cases for torch.onnx.ONNXProgram.optimize on NPU #43587
木路折创建于 8月2日
test: Add validation cases for torch.onnx.ONNXProgram.optimize on NPU #43587
已开启
共 1 个文件变更+47-1
| @@ -16,9 +16,11 @@ | |||
| 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 sufficient and direct API validations for some APIs, 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.model_proto (extendable). | 19 | +2. This file validates torch.onnx.ONNXProgram.model_proto, torch.onnx.ONNXProgram.optimize (extendable). |
| 20 | """ | 20 | """ |
| 21 | 21 | ||
| 22 | +from __future__ import annotations | ||
| 23 | + | ||
| 22 | import torch | 24 | import torch |
| 23 | from torch.testing._internal import common_utils | 25 | from torch.testing._internal import common_utils |
| 24 | 26 | ||
| @@ -83,5 +85,49 @@ class TestONNXProgramModelProto(common_utils.TestCase): | |||
| 83 | self.assertEqual(len(proto.graph.output), 2) | 85 | self.assertEqual(len(proto.graph.output), 2) |
| 84 | 86 | ||
| 85 | 87 | ||
| 88 | +class ONNXProgramOptimizeTest(common_utils.TestCase): | ||
| 89 | + """Tests for ONNXProgram.optimize method.""" | ||
| 90 | + | ||
| 91 | + def _create_onnx_program(self) -> torch.onnx.ONNXProgram: | ||
| 92 | + """Helper to create an unoptimized ONNXProgram from a simple model.""" | ||
| 93 | + class Model(torch.nn.Module): | ||
| 94 | + def forward(self, x): | ||
| 95 | + return (x + 1) * 2 | ||
| 96 | + | ||
| 97 | + x = torch.randn(3).to(device_type) | ||
| 98 | + onnx_program = torch.onnx.export( | ||
| 99 | + Model().eval(), (x,), dynamo=True, optimize=False, verbose=False | ||
| 100 | + ) | ||
| 101 | + return onnx_program | ||
| 102 | + | ||
| 103 | + def test_optimize_returns_none(self): | ||
| 104 | + """optimize() should return None.""" | ||
| 105 | + onnx_program = self._create_onnx_program() | ||
| 106 | + result = onnx_program.optimize() | ||
| 107 | + self.assertIsNone(result) | ||
| 108 | + | ||
| 109 | + def test_optimize_model_valid(self): | ||
| 110 | + """The model should be valid after optimize().""" | ||
| 111 | + onnx_program = self._create_onnx_program() | ||
| 112 | + onnx_program.optimize() | ||
| 113 | + self.assertIsNotNone(onnx_program.model) | ||
| 114 | + self.assertIsNotNone(onnx_program.model.graph) | ||
| 115 | + | ||
| 116 | + def test_optimize_idempotent(self): | ||
| 117 | + """Calling optimize() twice should not raise.""" | ||
| 118 | + onnx_program = self._create_onnx_program() | ||
| 119 | + onnx_program.optimize() | ||
| 120 | + onnx_program.optimize() | ||
| 121 | + | ||
| 122 | + def test_optimize_preserves_io_count(self): | ||
| 123 | + """The optimized model should preserve graph input/output count.""" | ||
| 124 | + onnx_program = self._create_onnx_program() | ||
| 125 | + num_inputs = len(onnx_program.model.graph.inputs) | ||
| 126 | + num_outputs = len(onnx_program.model.graph.outputs) | ||
| 127 | + onnx_program.optimize() | ||
| 128 | + self.assertEqual(len(onnx_program.model.graph.inputs), num_inputs) | ||
| 129 | + self.assertEqual(len(onnx_program.model.graph.outputs), num_outputs) | ||
| 130 | + | ||
| 131 | + | ||
| 86 | if __name__ == "__main__": | 132 | if __name__ == "__main__": |
| 87 | common_utils.run_tests() | 133 | common_utils.run_tests() |