已合并
fix(op-plugin): align native_batch_norm infer save tensors with meta/CUDA shapes. #4972
Margaret_wangrui创建于 5月18日
fix(op-plugin): align native_batch_norm infer save tensors with meta/CUDA shapes. #4972
已合并
共 3 个文件变更+62-8
| @@ -3653,10 +3653,10 @@ official: | |||
| 3653 | size: input | 3653 | size: input |
| 3654 | dtype: input | 3654 | dtype: input |
| 3655 | out1: | 3655 | out1: |
| 3656 | - size: 'training ? c10::SmallVector<int64_t, op_infer::SIZE>{input.size(1)} : c10::SmallVector<int64_t, op_infer::SIZE>{0}' | 3656 | + size: c10::SmallVector<int64_t, op_infer::SIZE>{input.size(1)} |
| 3657 | dtype: 'training ? at::kFloat : input.scalar_type()' | 3657 | dtype: 'training ? at::kFloat : input.scalar_type()' |
| 3658 | out2: | 3658 | out2: |
| 3659 | - size: 'training ? c10::SmallVector<int64_t, op_infer::SIZE>{input.size(1)} : c10::SmallVector<int64_t, op_infer::SIZE>{0}' | 3659 | + size: c10::SmallVector<int64_t, op_infer::SIZE>{input.size(1)} |
| 3660 | dtype: 'training ? at::kFloat : input.scalar_type()' | 3660 | dtype: 'training ? at::kFloat : input.scalar_type()' |
| 3661 | exec: aclnnBatchNorm | 3661 | exec: aclnnBatchNorm |
| 3662 | 3662 | ||
| @@ -214,8 +214,11 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> native_batch_norm( | |||
| 214 | npu_preparation::apply_tensor( | 214 | npu_preparation::apply_tensor( |
| 215 | running_var_tensor.sizes(), running_var_tensor.options().dtype(at::kFloat), self); | 215 | running_var_tensor.sizes(), running_var_tensor.options().dtype(at::kFloat), self); |
| 216 | } else { | 216 | } else { |
| 217 | - save_mean = at::empty({0}, self.options()); | 217 | + // Inference: BNInfer does not fill batch_mean/batch_variance outputs; shapes must still match |
| 218 | - save_invstd = at::empty({0}, self.options()); | 218 | + // PyTorch meta and CUDA ([num_features]), not empty ([0]), so run_meta_crossref / downstream |
| 219 | + // dispatch_meta_outplace_native_batch_norm_* tests agree with real PrivateUse1 outputs. | ||
| 220 | + save_mean = at::empty({dim_c}, self.options()); | ||
| 221 | + save_invstd = at::empty({dim_c}, self.options()); | ||
| 219 | } | 222 | } |
| 220 | 223 | ||
| 221 | return acl_op::native_batch_norm_out(self, weight_opt, bias_opt, | 224 | return acl_op::native_batch_norm_out(self, weight_opt, bias_opt, |
| @@ -1,11 +1,9 @@ | |||
| 1 | -import torch | 1 | +import unittest |
| 2 | -import numpy as np | ||
| 3 | 2 | ||
| 3 | +import torch | ||
| 4 | import torch_npu | 4 | import torch_npu |
| 5 | 5 | ||
| 6 | from torch_npu.testing.testcase import TestCase, run_tests | 6 | from torch_npu.testing.testcase import TestCase, run_tests |
| 7 | -from torch_npu.testing.common_utils import create_common_tensor | ||
| 8 | - | ||
| 9 | 7 | ||
| 10 | class TestNativeBatchNormLegit(TestCase): | 8 | class TestNativeBatchNormLegit(TestCase): |
| 11 | # pylint:disable = huawei-too-many-arguments | 9 | # pylint:disable = huawei-too-many-arguments |
| @@ -35,5 +33,58 @@ class TestNativeBatchNormLegit(TestCase): | |||
| 35 | self.assertRtolEqual(cpu_output[1], npu_output[1]) | 33 | self.assertRtolEqual(cpu_output[1], npu_output[1]) |
| 36 | 34 | ||
| 37 | 35 | ||
| 36 | +class TestNativeBatchNormInferSaveShapes(TestCase): | ||
| 37 | + """NPU inference save_mean/save_invstd match ATen meta ([num_features]); CPU inference uses empty ([0]).""" | ||
| 38 | + | ||
| 39 | + def test_infer_save_tensors_match_num_features_meta_crossref(self): | ||
| 40 | + if not torch.npu.is_available(): | ||
| 41 | + raise unittest.SkipTest("NPU not available") | ||
| 42 | + | ||
| 43 | + device = "npu" | ||
| 44 | + n, c = 3, 2 | ||
| 45 | + x = torch.randn(n, c, 4, dtype=torch.float32, device=device) | ||
| 46 | + weight = torch.randn(c, dtype=torch.float32, device=device) | ||
| 47 | + bias = torch.randn(c, dtype=torch.float32, device=device) | ||
| 48 | + running_mean = torch.randn(c, dtype=torch.float32, device=device) | ||
| 49 | + running_var = torch.randn(c, dtype=torch.float32, device=device).abs().add_(1e-3) | ||
| 50 | + | ||
| 51 | + out_npu, sm_npu, sis_npu = torch.ops.aten.native_batch_norm.default( | ||
| 52 | + x, weight, bias, running_mean, running_var, False, -1.2, 1e-5 | ||
| 53 | + ) | ||
| 54 | + | ||
| 55 | + self.assertEqual(out_npu.shape, x.shape) | ||
| 56 | + self.assertEqual(sm_npu.shape, (c,)) | ||
| 57 | + self.assertEqual(sis_npu.shape, (c,)) | ||
| 58 | + | ||
| 59 | + x_cpu = x.cpu() | ||
| 60 | + out_cpu, sm_cpu, sis_cpu = torch.ops.aten.native_batch_norm.default( | ||
| 61 | + x_cpu, | ||
| 62 | + weight.cpu(), | ||
| 63 | + bias.cpu(), | ||
| 64 | + running_mean.cpu(), | ||
| 65 | + running_var.cpu(), | ||
| 66 | + False, | ||
| 67 | + -1.2, | ||
| 68 | + 1e-5, | ||
| 69 | + ) | ||
| 70 | + # CPU inference returns empty save_mean / save_invstd ([0]); NPU matches ATen meta ([num_features]). | ||
| 71 | + self.assertEqual(tuple(sm_cpu.shape), (0,)) | ||
| 72 | + self.assertEqual(tuple(sis_cpu.shape), (0,)) | ||
| 73 | + | ||
| 74 | + meta_args = ( | ||
| 75 | + torch.empty(x.shape, device="meta"), | ||
| 76 | + torch.empty((c,), device="meta"), | ||
| 77 | + torch.empty((c,), device="meta"), | ||
| 78 | + torch.empty((c,), device="meta"), | ||
| 79 | + torch.empty((c,), device="meta"), | ||
| 80 | + False, | ||
| 81 | + -1.2, | ||
| 82 | + 1e-5, | ||
| 83 | + ) | ||
| 84 | + _, sm_meta, sis_meta = torch.ops.aten.native_batch_norm.default(*meta_args) | ||
| 85 | + self.assertEqual(tuple(sm_meta.shape), (c,)) | ||
| 86 | + self.assertEqual(tuple(sis_meta.shape), (c,)) | ||
| 87 | + | ||
| 88 | + | ||
| 38 | if __name__ == "__main__": | 89 | if __name__ == "__main__": |
| 39 | run_tests() | 90 | run_tests() |