Skip to content

Commit c4a40a0

Browse files
youge325youge325
committed
[Cpp API Compatibility] Add stft compatibility tests
1 parent a0e297e commit c4a40a0

1 file changed

Lines changed: 310 additions & 0 deletions

File tree

‎test/ATen/ops/StftTest.cpp‎

Lines changed: 310 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,310 @@
1+
#include <ATen/ATen.h>
2+
#include <ATen/core/Tensor.h>
3+
#include <ATen/ops/ones.h>
4+
#include <ATen/ops/stft.h>
5+
#include <gtest/gtest.h>
6+
7+
#include <cmath>
8+
#include <complex>
9+
#include <optional>
10+
#include <string>
11+
#include <vector>
12+
13+
#include "src/file_manager.h"
14+
15+
extern paddle_api_test::ThreadSafeParam g_custom_param;
16+
17+
namespace at {
18+
namespace test {
19+
20+
using paddle_api_test::FileManerger;
21+
using paddle_api_test::ThreadSafeParam;
22+
23+
// 写出 stft 结果: shape 信息 + dtype + 各 frame 的 DC bin 实部
24+
// 注: stft 的数值结果在 Paddle 与 PyTorch 之间可能因 padding/FFT
25+
// 实现差异而不同， 因此只比较 shape/dtype 和每帧 DC 分量(用于验证 FFT
26+
// 确实被执行)
27+
static void write_stft_result_to_file(FileManerger* file,
28+
const at::Tensor& result) {
29+
at::Tensor contig = result.contiguous();
30+
*file << std::to_string(contig.dim()) << " ";
31+
*file << std::to_string(contig.numel()) << " ";
32+
for (int64_t i = 0; i < contig.dim(); ++i) {
33+
*file << std::to_string(contig.sizes()[i]) << " ";
34+
}
35+
// 写出 dtype
36+
*file << std::to_string(static_cast<int>(contig.scalar_type())) << " ";
37+
38+
// 只写出每帧 DC 分量(bin 0)的实部，作为 FFT 已执行的验证
39+
int64_t n_frames = contig.sizes()[contig.dim() - 1];
40+
int64_t freq_bins = contig.sizes()[contig.dim() - 2];
41+
for (int64_t f = 0; f < n_frames; ++f) {
42+
if (contig.scalar_type() == at::kComplexFloat) {
43+
std::complex<float>* data =
44+
reinterpret_cast<std::complex<float>*>(contig.data_ptr());
45+
*file << std::to_string(data[f * freq_bins].real()) << " ";
46+
} else if (contig.scalar_type() == at::kComplexDouble) {
47+
std::complex<double>* data =
48+
reinterpret_cast<std::complex<double>*>(contig.data_ptr());
49+
*file << std::to_string(data[f * freq_bins].real()) << " ";
50+
} else if (contig.scalar_type() == at::kFloat) {
51+
// return_complex=false: shape [..., freq, frames, 2], last dim = [real,
52+
// imag]
53+
float* data = reinterpret_cast<float*>(contig.data_ptr());
54+
int64_t stride = freq_bins * 2;
55+
*file << std::to_string(data[f * stride]) << " ";
56+
} else if (contig.scalar_type() == at::kDouble) {
57+
double* data = reinterpret_cast<double*>(contig.data_ptr());
58+
int64_t stride = freq_bins * 2;
59+
*file << std::to_string(data[f * stride]) << " ";
60+
} else {
61+
*file << "unsupported_dtype ";
62+
}
63+
}
64+
}
65+
66+
class StftTest : public ::testing::Test {
67+
protected:
68+
void SetUp() override {
69+
// 基准 2D tensor: [batch=1, time=16], float32
70+
test_tensor = at::ones({1, 16}, at::kFloat);
71+
}
72+
at::Tensor test_tensor;
73+
};
74+
75+
// ========== 基础功能 ==========
76+
77+
TEST_F(StftTest, BasicStft) {
78+
auto file_name = g_custom_param.get();
79+
FileManerger file(file_name);
80+
file.createFile();
81+
file << "BasicStft ";
82+
at::Tensor result = test_tensor.stft(
83+
/*n_fft=*/8,
84+
/*hop_length=*/4,
85+
/*win_length=*/8,
86+
/*window=*/::std::nullopt,
87+
/*normalized=*/false,
88+
/*onesided=*/::std::nullopt,
89+
/*return_complex=*/true);
90+
write_stft_result_to_file(&file, result);
91+
file << "\n";
92+
file.saveFile();
93+
}
94+
95+
TEST_F(StftTest, StftWithWindow) {
96+
auto file_name = g_custom_param.get();
97+
FileManerger file(file_name);
98+
file.openAppend();
99+
file << "StftWithWindow ";
100+
at::Tensor window = at::ones({8}, at::kFloat);
101+
at::Tensor result = test_tensor.stft(
102+
/*n_fft=*/8,
103+
/*hop_length=*/4,
104+
/*win_length=*/8,
105+
/*window=*/window,
106+
/*normalized=*/false,
107+
/*onesided=*/true,
108+
/*return_complex=*/true);
109+
write_stft_result_to_file(&file, result);
110+
file << "\n";
111+
file.saveFile();
112+
}
113+
114+
TEST_F(StftTest, StftNormalized) {
115+
auto file_name = g_custom_param.get();
116+
FileManerger file(file_name);
117+
file.openAppend();
118+
file << "StftNormalized ";
119+
at::Tensor result = test_tensor.stft(
120+
/*n_fft=*/8,
121+
/*hop_length=*/4,
122+
/*win_length=*/8,
123+
/*window=*/::std::nullopt,
124+
/*normalized=*/true,
125+
/*onesided=*/true,
126+
/*return_complex=*/true);
127+
write_stft_result_to_file(&file, result);
128+
file << "\n";
129+
file.saveFile();
130+
}
131+
132+
// ========== Shape 覆盖 ==========
133+
134+
TEST_F(StftTest, SmallShape) {
135+
auto file_name = g_custom_param.get();
136+
FileManerger file(file_name);
137+
file.openAppend();
138+
file << "SmallShape ";
139+
at::Tensor t = at::ones({1, 16}, at::kFloat);
140+
at::Tensor result = t.stft(/*n_fft=*/8,
141+
/*hop_length=*/4,
142+
/*win_length=*/8,
143+
/*window=*/::std::nullopt,
144+
/*normalized=*/false,
145+
/*onesided=*/true,
146+
/*return_complex=*/true);
147+
write_stft_result_to_file(&file, result);
148+
file << "\n";
149+
file.saveFile();
150+
}
151+
152+
TEST_F(StftTest, LargeShape) {
153+
auto file_name = g_custom_param.get();
154+
FileManerger file(file_name);
155+
file.openAppend();
156+
file << "LargeShape ";
157+
at::Tensor t = at::ones({2, 64}, at::kFloat);
158+
at::Tensor result = t.stft(/*n_fft=*/16,
159+
/*hop_length=*/8,
160+
/*win_length=*/16,
161+
/*window=*/::std::nullopt,
162+
/*normalized=*/false,
163+
/*onesided=*/true,
164+
/*return_complex=*/true);
165+
write_stft_result_to_file(&file, result);
166+
file << "\n";
167+
file.saveFile();
168+
}
169+
170+
TEST_F(StftTest, OneDInput) {
171+
auto file_name = g_custom_param.get();
172+
FileManerger file(file_name);
173+
file.openAppend();
174+
file << "OneDInput ";
175+
at::Tensor t = at::ones({16}, at::kFloat);
176+
at::Tensor result = t.stft(/*n_fft=*/8,
177+
/*hop_length=*/4,
178+
/*win_length=*/8,
179+
/*window=*/::std::nullopt,
180+
/*normalized=*/false,
181+
/*onesided=*/true,
182+
/*return_complex=*/true);
183+
write_stft_result_to_file(&file, result);
184+
file << "\n";
185+
file.saveFile();
186+
}
187+
188+
// ========== Dtype 覆盖 ==========
189+
190+
TEST_F(StftTest, Float64Dtype) {
191+
auto file_name = g_custom_param.get();
192+
FileManerger file(file_name);
193+
file.openAppend();
194+
file << "Float64Dtype ";
195+
at::Tensor t = at::ones({1, 16}, at::kDouble);
196+
at::Tensor result = t.stft(/*n_fft=*/8,
197+
/*hop_length=*/4,
198+
/*win_length=*/8,
199+
/*window=*/::std::nullopt,
200+
/*normalized=*/false,
201+
/*onesided=*/true,
202+
/*return_complex=*/true);
203+
write_stft_result_to_file(&file, result);
204+
file << "\n";
205+
file.saveFile();
206+
}
207+
208+
// ========== API 变体 ==========
209+
210+
TEST_F(StftTest, DifferentHopLength) {
211+
auto file_name = g_custom_param.get();
212+
FileManerger file(file_name);
213+
file.openAppend();
214+
file << "DifferentHopLength ";
215+
at::Tensor t = at::ones({1, 32}, at::kFloat);
216+
at::Tensor result = t.stft(/*n_fft=*/16,
217+
/*hop_length=*/8,
218+
/*win_length=*/16,
219+
/*window=*/::std::nullopt,
220+
/*normalized=*/false,
221+
/*onesided=*/true,
222+
/*return_complex=*/true);
223+
write_stft_result_to_file(&file, result);
224+
file << "\n";
225+
file.saveFile();
226+
}
227+
228+
TEST_F(StftTest, ReturnComplexFalse) {
229+
auto file_name = g_custom_param.get();
230+
FileManerger file(file_name);
231+
file.openAppend();
232+
file << "ReturnComplexFalse ";
233+
at::Tensor result = test_tensor.stft(
234+
/*n_fft=*/8,
235+
/*hop_length=*/4,
236+
/*win_length=*/8,
237+
/*window=*/::std::nullopt,
238+
/*normalized=*/false,
239+
/*onesided=*/true,
240+
/*return_complex=*/false);
241+
write_stft_result_to_file(&file, result);
242+
file << "\n";
243+
file.saveFile();
244+
}
245+
246+
TEST_F(StftTest, PyTorchStyleCenterFalseOverload) {
247+
auto file_name = g_custom_param.get();
248+
FileManerger file(file_name);
249+
file.openAppend();
250+
file << "PyTorchStyleCenterFalseOverload ";
251+
at::Tensor result = test_tensor.stft(
252+
/*n_fft=*/8,
253+
/*hop_length=*/4,
254+
/*win_length=*/8,
255+
/*window=*/::std::nullopt,
256+
/*center=*/false,
257+
/*pad_mode=*/"reflect",
258+
/*normalized=*/true,
259+
/*onesided=*/true,
260+
/*return_complex=*/true);
261+
write_stft_result_to_file(&file, result);
262+
file << "\n";
263+
file.saveFile();
264+
}
265+
266+
TEST_F(StftTest, ReturnComplexUnspecifiedRealInputThrows) {
267+
auto file_name = g_custom_param.get();
268+
FileManerger file(file_name);
269+
file.openAppend();
270+
file << "ReturnComplexUnspecifiedRealInputThrows ";
271+
bool thrown = false;
272+
try {
273+
(void)test_tensor.stft(/*n_fft=*/8,
274+
/*hop_length=*/4,
275+
/*win_length=*/8,
276+
/*window=*/::std::nullopt,
277+
/*normalized=*/false,
278+
/*onesided=*/true);
279+
} catch (const std::exception&) {
280+
thrown = true;
281+
file << "exception ";
282+
}
283+
if (!thrown) {
284+
file << "no_exception ";
285+
}
286+
file << "\n";
287+
file.saveFile();
288+
}
289+
290+
TEST_F(StftTest, WinLengthSmallerThanNFFT) {
291+
auto file_name = g_custom_param.get();
292+
FileManerger file(file_name);
293+
file.openAppend();
294+
file << "WinLengthSmallerThanNFFT ";
295+
at::Tensor window = at::ones({4}, at::kFloat);
296+
at::Tensor result = test_tensor.stft(
297+
/*n_fft=*/8,
298+
/*hop_length=*/4,
299+
/*win_length=*/4,
300+
/*window=*/window,
301+
/*normalized=*/false,
302+
/*onesided=*/true,
303+
/*return_complex=*/true);
304+
write_stft_result_to_file(&file, result);
305+
file << "\n";
306+
file.saveFile();
307+
}
308+
309+
} // namespace test
310+
} // namespace at

0 commit comments

Comments
 (0)