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
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
| """
频域特征提取模块
使用STFT提取频谱特征，CNN提取频域模式
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple
class SpectralFeatureExtractor(nn.Module):
"""
频域特征提取器
步骤：
1. STFT变换
2. 功率谱密度计算
3. CNN提取频域模式
输入: (batch, channels, time_steps)
输出: (batch, features, freq_bins)
"""
def __init__(
self,
sample_rate: int = 256,
n_fft: int = 64,
hop_length: int = 16,
n_mels: int = 32,
embed_dim: int = 64
):
super().__init__()
self.sample_rate = sample_rate
self.n_fft = n_fft
self.hop_length = hop_length
self.n_mels = n_mels
self.freq_cnn = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=(3, 3), padding=1),
nn.BatchNorm2d(32),
nn.GELU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, kernel_size=(3, 3), padding=1),
nn.BatchNorm2d(64),
nn.GELU(),
nn.MaxPool2d(2),
nn.Conv2d(64, embed_dim, kernel_size=(3, 3), padding=1),
nn.BatchNorm2d(embed_dim),
nn.GELU(),
)
freq_bins = n_fft // 2 + 1
self.freq_attention = nn.Sequential(
nn.AdaptiveAvgPool2d((freq_bins, 1)),
nn.Conv2d(embed_dim, embed_dim // 4, kernel_size=1),
nn.GELU(),
nn.Conv2d(embed_dim // 4, embed_dim, kernel_size=1),
nn.Sigmoid()
)
self.out_dim = embed_dim
def compute_spectrogram(
self,
x: torch.Tensor
) -> torch.Tensor:
"""
计算频谱图
Args:
x: (batch, channels, time_steps)
Returns:
spec: (batch, channels, freq, time)
"""
batch, channels, time_steps = x.shape
specs = []
for ch in range(channels):
window = torch.hann_window(self.n_fft, device=x.device)
stft = torch.stft(
x[:, ch, :],
n_fft=self.n_fft,
hop_length=self.hop_length,
window=window,
return_complex=True
)
power = torch.abs(stft) ** 2
specs.append(power)
spec = torch.stack(specs, dim=1)
spec = spec.mean(dim=1, keepdim=True)
spec = torch.log(spec + 1e-8)
return spec
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: (batch, channels, time_steps)
Returns:
features: (batch, out_dim, freq_bins)
"""
spec = self.compute_spectrogram(x)
features = self.freq_cnn(spec)
attention = self.freq_attention(features)
features = features * attention
features = features.mean(dim=-1)
return features
if __name__ == "__main__":
batch_size = 8
n_channels = 14
time_steps = 256
x = torch.randn(batch_size, n_channels, time_steps)
model = SpectralFeatureExtractor(
sample_rate=256,
n_fft=64,
hop_length=16
)
features = model(x)
print("=== 频域特征提取测试 ===")
print(f"输入形状: {x.shape}")
print(f"输出形状: {features.shape}")
print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")
|