STFT/iSTFT without unfold/complex ops, using conv1d and conv_transpose1d. - forward STFT => Real-part conv1d + Imag-part conv1d - inverse STFT => Real-part conv_transpose1d + Imag-part conv_transpose1d + sum - avoids F.unfold, so easier to export to ONNX - uses replicate or con
| 5 | import torch.nn.functional as F |
| 6 | |
| 7 | class CustomSTFT(nn.Module): |
| 8 | """ |
| 9 | STFT/iSTFT without unfold/complex ops, using conv1d and conv_transpose1d. |
| 10 | |
| 11 | - forward STFT => Real-part conv1d + Imag-part conv1d |
| 12 | - inverse STFT => Real-part conv_transpose1d + Imag-part conv_transpose1d + sum |
| 13 | - avoids F.unfold, so easier to export to ONNX |
| 14 | - uses replicate or constant padding for 'center=True' to approximate 'reflect' |
| 15 | (reflect is not supported for dynamic shapes in ONNX) |
| 16 | """ |
| 17 | |
| 18 | def __init__( |
| 19 | self, |
| 20 | filter_length=800, |
| 21 | hop_length=200, |
| 22 | win_length=800, |
| 23 | window="hann", |
| 24 | center=True, |
| 25 | pad_mode="replicate", # or 'constant' |
| 26 | ): |
| 27 | super().__init__() |
| 28 | self.filter_length = filter_length |
| 29 | self.hop_length = hop_length |
| 30 | self.win_length = win_length |
| 31 | self.n_fft = filter_length |
| 32 | self.center = center |
| 33 | self.pad_mode = pad_mode |
| 34 | |
| 35 | # Number of frequency bins for real-valued STFT with onesided=True |
| 36 | self.freq_bins = self.n_fft // 2 + 1 |
| 37 | |
| 38 | # Build window |
| 39 | assert window == 'hann', window |
| 40 | window_tensor = torch.hann_window(win_length, periodic=True, dtype=torch.float32) |
| 41 | if self.win_length < self.n_fft: |
| 42 | # Zero-pad up to n_fft |
| 43 | extra = self.n_fft - self.win_length |
| 44 | window_tensor = F.pad(window_tensor, (0, extra)) |
| 45 | elif self.win_length > self.n_fft: |
| 46 | window_tensor = window_tensor[: self.n_fft] |
| 47 | self.register_buffer("window", window_tensor) |
| 48 | |
| 49 | # Precompute forward DFT (real, imag) |
| 50 | # PyTorch stft uses e^{-j 2 pi k n / N} => real=cos(...), imag=-sin(...) |
| 51 | n = np.arange(self.n_fft) |
| 52 | k = np.arange(self.freq_bins) |
| 53 | angle = 2 * np.pi * np.outer(k, n) / self.n_fft # shape (freq_bins, n_fft) |
| 54 | dft_real = np.cos(angle) |
| 55 | dft_imag = -np.sin(angle) # note negative sign |
| 56 | |
| 57 | # Combine window and dft => shape (freq_bins, filter_length) |
| 58 | # We'll make 2 conv weight tensors of shape (freq_bins, 1, filter_length). |
| 59 | forward_window = window_tensor.numpy() # shape (n_fft,) |
| 60 | forward_real = dft_real * forward_window # (freq_bins, n_fft) |
| 61 | forward_imag = dft_imag * forward_window |
| 62 | |
| 63 | # Convert to PyTorch |
| 64 | forward_real_torch = torch.from_numpy(forward_real).float() |