Skip to content

Commit 2b47a0a

Browse files
youge325youge325
committed
test(ATen): add stft alignment coverage
1 parent c29c30b commit 2b47a0a

1 file changed

Lines changed: 264 additions & 0 deletions

File tree

‎test/ATen/ops/StftTest.cpp‎

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

0 commit comments

Comments
 (0)