已合并
fix torch.stft bug on ouput shape when input is empty tensor #5332
XianglongZeng创建于 7月1日
fix torch.stft bug on ouput shape when input is empty tensor #5332
已合并
共 2 个文件变更+49-10
| @@ -56,7 +56,7 @@ at::ScalarType get_output_type(bool return_complex, at::ScalarType input_type) | |||
| 56 | return output_type; | 56 | return output_type; |
| 57 | } | 57 | } |
| 58 | 58 | ||
| 59 | -c10::SmallVector<int64_t, SIZE> get_output_size(bool return_complex, int64_t batch, int64_t frames, int64_t n) | 59 | +c10::SmallVector<int64_t, SIZE> get_output_size(bool return_complex, bool has_batch, int64_t batch, int64_t frames, int64_t n) |
| 60 | { | 60 | { |
| 61 | c10::SmallVector<int64_t, SIZE> output_size; | 61 | c10::SmallVector<int64_t, SIZE> output_size; |
| 62 | c10::SmallVector<int64_t, SIZE> output_complex_with_batch = {batch, n, frames}; | 62 | c10::SmallVector<int64_t, SIZE> output_complex_with_batch = {batch, n, frames}; |
| @@ -65,9 +65,9 @@ c10::SmallVector<int64_t, SIZE> get_output_size(bool return_complex, int64_t bat | |||
| 65 | c10::SmallVector<int64_t, SIZE> output_real = {n, frames, 2}; | 65 | c10::SmallVector<int64_t, SIZE> output_real = {n, frames, 2}; |
| 66 | 66 | ||
| 67 | if (return_complex) { | 67 | if (return_complex) { |
| 68 | - output_size = batch > 0 ? output_complex_with_batch : output_complex; | 68 | + output_size = has_batch ? output_complex_with_batch : output_complex; |
| 69 | } else { | 69 | } else { |
| 70 | - output_size = batch > 0 ? output_real_with_batch : output_real; | 70 | + output_size = has_batch ? output_real_with_batch : output_real; |
| 71 | } | 71 | } |
| 72 | return output_size; | 72 | return output_size; |
| 73 | } | 73 | } |
| @@ -83,7 +83,7 @@ bool is_aclstft_supported(at::Tensor const &self, | |||
| 83 | if (self.scalar_type() == at::ScalarType::Double || self.scalar_type() == at::ScalarType::ComplexDouble) { | 83 | if (self.scalar_type() == at::ScalarType::Double || self.scalar_type() == at::ScalarType::ComplexDouble) { |
| 84 | return true; | 84 | return true; |
| 85 | } | 85 | } |
| 86 | - | 86 | + |
| 87 | bool res = false; | 87 | bool res = false; |
| 88 | int64_t batch = self.dim() == 2 ? self.size(0) : 0; | 88 | int64_t batch = self.dim() == 2 ? self.size(0) : 0; |
| 89 | int64_t len = self.dim() == 2 ? self.size(1) : self.size(0); | 89 | int64_t len = self.dim() == 2 ? self.size(1) : self.size(0); |
| @@ -221,8 +221,9 @@ at::Tensor stft(at::Tensor const &self, | |||
| 221 | return_complex.value() : | 221 | return_complex.value() : |
| 222 | self.is_complex() || (window.has_value() && window.value().is_complex()); | 222 | self.is_complex() || (window.has_value() && window.value().is_complex()); |
| 223 | 223 | ||
| 224 | - int64_t batch = self.dim() == 2 ? self.size(0) : 0; | 224 | + bool has_batch = self.dim() == 2; |
| 225 | - int64_t len = self.dim() == 2 ? self.size(1) : self.size(0); | 225 | + int64_t batch = has_batch ? self.size(0) : 0; |
| 226 | + int64_t len = has_batch ? self.size(1) : self.size(0); | ||
| 226 | int64_t frames = (len - n_fft) / hop_length_value + 1; | 227 | int64_t frames = (len - n_fft) / hop_length_value + 1; |
| 227 | int64_t n = onesided_value ? n_fft / 2 + 1 : n_fft; | 228 | int64_t n = onesided_value ? n_fft / 2 + 1 : n_fft; |
| 228 | at::ScalarType output_type = get_output_type(return_complex_value, self.scalar_type()); | 229 | at::ScalarType output_type = get_output_type(return_complex_value, self.scalar_type()); |
| @@ -230,11 +231,21 @@ at::Tensor stft(at::Tensor const &self, | |||
| 230 | output_type == at::ScalarType::ComplexFloat || output_type == at::ScalarType::ComplexDouble, | 231 | output_type == at::ScalarType::ComplexFloat || output_type == at::ScalarType::ComplexDouble, |
| 231 | "output type should be float, double, complex<float> or complex<double>", OPS_ERROR(ErrCode::TYPE)); | 232 | "output type should be float, double, complex<float> or complex<double>", OPS_ERROR(ErrCode::TYPE)); |
| 232 | 233 | ||
| 234 | + c10::SmallVector<int64_t, SIZE> output_size = get_output_size(return_complex_value, has_batch, batch, frames, n); | ||
| 235 | + | ||
| 236 | + // Neither aclStft nor the _stft composite path handles empty input correctly, | ||
| 237 | + // return an empty output with the expected shape directly. | ||
| 238 | + if (self.numel() == 0) { | ||
| 239 | + for (auto& d : output_size) { | ||
| 240 | + if (d < 0) { d = 0; } | ||
| 241 | + } | ||
| 242 | + return at::empty(output_size, self.options().dtype(output_type)); | ||
| 243 | + } | ||
| 244 | + | ||
| 233 | if (!is_aclstft_supported(self, window_value, n_fft, hop_length_value, normalized, onesided_value, return_complex_value)) { | 245 | if (!is_aclstft_supported(self, window_value, n_fft, hop_length_value, normalized, onesided_value, return_complex_value)) { |
| 234 | return _stft(self, n_fft, hop_length, win_length, window, normalized, onesided, return_complex, c10::nullopt); | 246 | return _stft(self, n_fft, hop_length, win_length, window, normalized, onesided, return_complex, c10::nullopt); |
| 235 | } | 247 | } |
| 236 | 248 | ||
| 237 | - c10::SmallVector<int64_t, SIZE> output_size = get_output_size(return_complex_value, batch, frames, n); | ||
| 238 | at::Tensor output = npu_preparation::apply_tensor_without_format(output_size, self.options().dtype(output_type)); | 249 | at::Tensor output = npu_preparation::apply_tensor_without_format(output_size, self.options().dtype(output_type)); |
| 239 | 250 | ||
| 240 | EXEC_NPU_CMD(aclStft, self, window_value, output, n_fft, hop_length_value, win_length_value, normalized, onesided_value, return_complex_value); | 251 | EXEC_NPU_CMD(aclStft, self, window_value, output, n_fft, hop_length_value, win_length_value, normalized, onesided_value, return_complex_value); |
| @@ -273,8 +284,9 @@ at::Tensor stft(at::Tensor const &self, | |||
| 273 | return_complex_opt.value() : | 284 | return_complex_opt.value() : |
| 274 | self.is_complex() || (window_opt.has_value() && window_opt.value().is_complex()); | 285 | self.is_complex() || (window_opt.has_value() && window_opt.value().is_complex()); |
| 275 | 286 | ||
| 276 | - int64_t batch = self.dim() == 2 ? self.size(0) : 0; | 287 | + bool has_batch = self.dim() == 2; |
| 277 | - int64_t len = self.dim() == 2 ? self.size(1) : self.size(0); | 288 | + int64_t batch = has_batch ? self.size(0) : 0; |
| 289 | + int64_t len = has_batch ? self.size(1) : self.size(0); | ||
| 278 | int64_t frames = (len - n_fft) / hop_length + 1; | 290 | int64_t frames = (len - n_fft) / hop_length + 1; |
| 279 | int64_t n = onesided == true ? n_fft / 2 + 1 : n_fft; | 291 | int64_t n = onesided == true ? n_fft / 2 + 1 : n_fft; |
| 280 | at::ScalarType output_type = get_output_type(return_complex, self.scalar_type()); | 292 | at::ScalarType output_type = get_output_type(return_complex, self.scalar_type()); |
| @@ -282,11 +294,21 @@ at::Tensor stft(at::Tensor const &self, | |||
| 282 | output_type == at::ScalarType::ComplexFloat || output_type == at::ScalarType::ComplexDouble, | 294 | output_type == at::ScalarType::ComplexFloat || output_type == at::ScalarType::ComplexDouble, |
| 283 | "output type should be float, double, complex<float> or complex<double>", OPS_ERROR(ErrCode::TYPE)); | 295 | "output type should be float, double, complex<float> or complex<double>", OPS_ERROR(ErrCode::TYPE)); |
| 284 | 296 | ||
| 297 | + c10::SmallVector<int64_t, SIZE> output_size = get_output_size(return_complex, has_batch, batch, frames, n); | ||
| 298 | + | ||
| 299 | + // Neither aclStft nor the _stft composite path handles empty input correctly, | ||
| 300 | + // return an empty output with the expected shape directly. | ||
| 301 | + if (self.numel() == 0) { | ||
| 302 | + for (auto& d : output_size) { | ||
| 303 | + if (d < 0) { d = 0; } | ||
| 304 | + } | ||
| 305 | + return at::empty(output_size, self.options().dtype(output_type)); | ||
| 306 | + } | ||
| 307 | + | ||
| 285 | if (!is_aclstft_supported(self, window, n_fft, hop_length, normalized, onesided, return_complex)) { | 308 | if (!is_aclstft_supported(self, window, n_fft, hop_length, normalized, onesided, return_complex)) { |
| 286 | return _stft(self, n_fft, hop_length_opt, win_length_opt, window_opt, normalized, onesided_opt, return_complex_opt, align_to_window); | 309 | return _stft(self, n_fft, hop_length_opt, win_length_opt, window_opt, normalized, onesided_opt, return_complex_opt, align_to_window); |
| 287 | } | 310 | } |
| 288 | 311 | ||
| 289 | - c10::SmallVector<int64_t, SIZE> output_size = get_output_size(return_complex, batch, frames, n); | ||
| 290 | at::Tensor output = npu_preparation::apply_tensor_without_format(output_size, self.options().dtype(output_type)); | 312 | at::Tensor output = npu_preparation::apply_tensor_without_format(output_size, self.options().dtype(output_type)); |
| 291 | 313 | ||
| 292 | EXEC_NPU_CMD(aclStft, self, window, output, n_fft, hop_length, win_length, normalized, onesided, return_complex); | 314 | EXEC_NPU_CMD(aclStft, self, window, output, n_fft, hop_length, win_length, normalized, onesided, return_complex); |
| @@ -23,5 +23,22 @@ class TestSTFT(TestCase): | |||
| 23 | 23 | ||
| 24 | self.assertRtolEqual(cpu_output, npu_output) | 24 | self.assertRtolEqual(cpu_output, npu_output) |
| 25 | 25 | ||
| 26 | + | ||
| 27 | + def test_stft_empty_input(self): | ||
| 28 | + input_tensor = torch.zeros(0, 19, dtype=torch.float32) | ||
| 29 | + window = torch.ones(14, dtype=torch.float32) | ||
| 30 | + res = torch.stft(input_tensor, 18, 8, 14, window=window, center=False, | ||
| 31 | + normalized=True, onesided=False, return_complex=False) | ||
| 32 | + | ||
| 33 | + input_tensor_npu = input_tensor.npu() | ||
| 34 | + window_npu = window.npu() | ||
| 35 | + res_npu = torch.stft(input_tensor_npu, 18, 8, 14, window=window_npu, center=False, | ||
| 36 | + normalized=True, onesided=False, return_complex=False) | ||
| 37 | + | ||
| 38 | + self.assertEqual(res.shape, res_npu.shape) | ||
| 39 | + self.assertEqual(res.shape, torch.Size([0, 18, 1, 2])) | ||
| 40 | + self.assertEqual(res.numel(), 0) | ||
| 41 | + self.assertEqual(res_npu.numel(), 0) | ||
| 42 | + | ||
| 26 | if __name__ == "__main__": | 43 | if __name__ == "__main__": |
| 27 | run_tests() | 44 | run_tests() |