已合并
test: add NPU coverage for Optimizer.zero_grad #43537
Jinfan Liu创建于 8月1日
test: add NPU coverage for Optimizer.zero_grad #43537
已合并
共 1 个文件变更+77-0
| @@ -0,0 +1,77 @@ | |||
| 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 | ||
| 13 | +# implied. | ||
| 14 | +# See the License for the specific language governing permissions and | ||
| 15 | +# limitations under the License. | ||
| 16 | + | ||
群 | |||
| 17 | +""" | ||
| 18 | +Add validation cases for torch.optim APIs on NPU: | ||
| 19 | +1. PyTorch community lacks sufficient and direct API validations for some APIs, so this file is added. | ||
| 20 | +2. This file validates torch.optim.Optimizer.zero_grad (extendable). | ||
| 21 | +""" | ||
| 22 | + | ||
| 23 | +import torch | ||
| 24 | +import torch_npu | ||
| 25 | + | ||
| 26 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +class TestOptimizerZeroGrad(TestCase): | ||
| 30 | + def _create_optimizer_with_gradient(self): | ||
| 31 | + parameter = torch.nn.Parameter(torch.tensor([2.0, -3.0], device="npu")) | ||
| 32 | + optimizer = torch.optim.SGD([parameter], lr=0.1) | ||
| 33 | + parameter.square().sum().backward() | ||
| 34 | + self.assertIsNotNone(parameter.grad) | ||
| 35 | + self.assertTrue(torch.count_nonzero(parameter.grad).item() > 0) | ||
| 36 | + return parameter, optimizer | ||
| 37 | + | ||
| 38 | + def _assert_gradient_is_none(self, set_to_none): | ||
| 39 | + parameter, optimizer = self._create_optimizer_with_gradient() | ||
| 40 | + optimizer.zero_grad(set_to_none) | ||
| 41 | + self.assertIsNone(parameter.grad) | ||
| 42 | + | ||
| 43 | + def _assert_gradient_is_zero(self, set_to_none): | ||
| 44 | + parameter, optimizer = self._create_optimizer_with_gradient() | ||
| 45 | + optimizer.zero_grad(set_to_none) | ||
| 46 | + self.assertIsNotNone(parameter.grad) | ||
| 47 | + self.assertTrue(torch.equal(parameter.grad, torch.zeros_like(parameter.grad))) | ||
| 48 | + | ||
| 49 | + def test_zero_grad_default_sets_gradient_to_none(self): | ||
| 50 | + parameter, optimizer = self._create_optimizer_with_gradient() | ||
| 51 | + | ||
| 52 | + optimizer.zero_grad() | ||
| 53 | + | ||
| 54 | + self.assertIsNone(parameter.grad) | ||
| 55 | + | ||
| 56 | + def test_zero_grad_boolean_parameter_values(self): | ||
| 57 | + self._assert_gradient_is_none(True) | ||
| 58 | + self._assert_gradient_is_zero(False) | ||
| 59 | + | ||
| 60 | + def test_zero_grad_non_boolean_truthy_and_falsy_values(self): | ||
| 61 | + # The upstream implementation uses Python truthiness for this argument. | ||
| 62 | + for set_to_none in (1, [1]): | ||
| 63 | + self._assert_gradient_is_none(set_to_none) | ||
| 64 | + for set_to_none in (0, [], None): | ||
| 65 | + self._assert_gradient_is_zero(set_to_none) | ||
| 66 | + | ||
| 67 | + def test_zero_grad_rejects_invalid_call_signatures(self): | ||
| 68 | + _, optimizer = self._create_optimizer_with_gradient() | ||
| 69 | + | ||
| 70 | + with self.assertRaises(TypeError): | ||
| 71 | + optimizer.zero_grad(True, False) | ||
| 72 | + with self.assertRaises(TypeError): | ||
| 73 | + optimizer.zero_grad(unexpected=True) | ||
| 74 | + | ||
| 75 | + | ||
| 76 | +if __name__ == "__main__": | ||
| 77 | + run_tests() | ||
缺少说明