已合并
Resize_ to support ncdhw to other dims #22678
wang-guangbin创建于 2025年7月4日
Resize_ to support ncdhw to other dims #22678
已合并
从refs/pull/22678/head合入到master
共 2 个文件变更+38-2
| @@ -0,0 +1,32 @@ | |||
| 1 | +import numpy as np | ||
| 2 | +import torch | ||
| 3 | +import torch_npu | ||
| 4 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 5 | + | ||
| 6 | + | ||
| 7 | +class TestResize(TestCase): | ||
| 8 | + | ||
| 9 | + def test_masked_select_out(self): | ||
| 10 | + | ||
| 11 | + input_data = torch.tensor([[[[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20], [21, 22, 23, 24, 25]]]]]], dtype=torch.float) | ||
| 12 | + mask = torch.tensor([True, False, True, False, True]) | ||
| 13 | + | ||
| 14 | + input_data_npu = input_data.npu() | ||
| 15 | + mask_npu = mask.npu() | ||
| 16 | + | ||
| 17 | + out_tensor = torch.empty((1, 1, 1, 1, 1), dtype=input_data.dtype) | ||
| 18 | + out_tensor_npu = out_tensor.npu() | ||
| 19 | + | ||
| 20 | + out_tensor_npu = out_tensor_npu.view(-1) | ||
| 21 | + out_tensor_npu = torch.masked_select(input_data_npu, mask_npu, out=out_tensor_npu) | ||
| 22 | + out_tensor = torch.masked_select(input_data, mask, out=out_tensor) | ||
| 23 | + self.assertRtolEqual(out_tensor_npu, out_tensor) | ||
| 24 | + | ||
| 25 | + def test_resize_ncdhw(self): | ||
| 26 | + out_tensor = torch.empty((1, 1, 1, 1, 1), dtype=torch.float16).npu() | ||
| 27 | + shape = [25] | ||
| 28 | + out_tensor.resize_(shape) | ||
| 29 | + self.assertEqual(shape, out_tensor.shape) | ||
| 30 | + | ||
| 31 | +if __name__ == "__main__": | ||
| 32 | + run_tests() | ||
| @@ -62,9 +62,13 @@ void StorageDescHelper::UpdateDesc(torch_npu::NPUStorageDesc &npuDesc, const c10 | |||
| 62 | } | 62 | } |
| 63 | } | 63 | } |
| 64 | npuDesc.base_strides_ = new_stride; | 64 | npuDesc.base_strides_ = new_stride; |
| 65 | - | ||
| 66 | // 更新物理内存信息 | 65 | // 更新物理内存信息 |
| 67 | - npuDesc.storage_sizes_ = FormatHelper::GetStorageSizes(npuDesc); | 66 | + int NCDHW_OR_NDHWC_DIM = 5; |
| 67 | + if ((npuDesc.npu_format_ == ACL_FORMAT_NCDHW || npuDesc.npu_format_ == ACL_FORMAT_NDHWC) && new_size.size() < NCDHW_OR_NDHWC_DIM) { | ||
| 68 | + npuDesc.storage_sizes_ = new_size; | ||
| 69 | + } else { | ||
| 70 | + npuDesc.storage_sizes_ = FormatHelper::GetStorageSizes(npuDesc); | ||
| 71 | + } | ||
| 68 | if (new_data_numel > new_shape_numel) { | 72 | if (new_data_numel > new_shape_numel) { |
| 69 | // Refresh format to base format only when flattening storage data | 73 | // Refresh format to base format only when flattening storage data |
| 70 | npuDesc.storage_sizes_ = new_size; | 74 | npuDesc.storage_sizes_ = new_size; |