| import torch |
| import torch.nn as nn |
| from torch.nn.modules.rnn import LSTM |
|
|
|
|
| class FeatureConversion(nn.Module): |
| """ |
| Integrates into the adjacent Dual-Path layer. |
| |
| Args: |
| channels (int): Number of input channels. |
| inverse (bool): If True, uses ifft; otherwise, uses rfft. |
| """ |
|
|
| def __init__(self, channels, inverse): |
| super().__init__() |
| self.inverse = inverse |
| self.channels = channels |
|
|
| def forward(self, x): |
| |
| if self.inverse: |
| x = x.float() |
| x_r = x[:, :self.channels // 2, :, :] |
| x_i = x[:, self.channels // 2:, :, :] |
| x = torch.complex(x_r, x_i) |
| x = torch.fft.irfft(x, dim=3, norm="ortho") |
| else: |
| x = x.float() |
| x = torch.fft.rfft(x, dim=3, norm="ortho") |
| x_real = x.real |
| x_imag = x.imag |
| x = torch.cat([x_real, x_imag], dim=1) |
| return x |
|
|
|
|
| class DualPathRNN(nn.Module): |
| """ |
| Dual-Path RNN in Separation Network. |
| |
| Args: |
| d_model (int): The number of expected features in the input (input_size). |
| expand (int): Expansion factor used to calculate the hidden_size of LSTM. |
| bidirectional (bool): If True, becomes a bidirectional LSTM. |
| """ |
|
|
| def __init__(self, d_model, expand, bidirectional=True): |
| super(DualPathRNN, self).__init__() |
|
|
| self.d_model = d_model |
| self.hidden_size = d_model * expand |
| self.bidirectional = bidirectional |
| |
| self.lstm_layers = nn.ModuleList([self._init_lstm_layer(self.d_model, self.hidden_size) for _ in range(2)]) |
| self.linear_layers = nn.ModuleList([nn.Linear(self.hidden_size * 2, self.d_model) for _ in range(2)]) |
| self.norm_layers = nn.ModuleList([nn.GroupNorm(1, d_model) for _ in range(2)]) |
|
|
| def _init_lstm_layer(self, d_model, hidden_size): |
| return LSTM(d_model, hidden_size, num_layers=1, bidirectional=self.bidirectional, batch_first=True) |
|
|
| def forward(self, x): |
| B, C, F, T = x.shape |
|
|
| |
| original_x = x |
| |
| x = self.norm_layers[0](x) |
| x = x.transpose(1, 3).contiguous().view(B * T, F, C) |
| x, _ = self.lstm_layers[0](x) |
| x = self.linear_layers[0](x) |
| x = x.view(B, T, F, C).transpose(1, 3) |
| x = x + original_x |
|
|
| original_x = x |
| |
| x = self.norm_layers[1](x) |
| x = x.transpose(1, 2).contiguous().view(B * F, C, T).transpose(1, 2) |
| x, _ = self.lstm_layers[1](x) |
| x = self.linear_layers[1](x) |
| x = x.transpose(1, 2).contiguous().view(B, F, C, T).transpose(1, 2) |
| x = x + original_x |
|
|
| return x |
|
|
|
|
| class SeparationNet(nn.Module): |
| """ |
| Implements a simplified Sparse Down-sample block in an encoder architecture. |
| |
| Args: |
| - channels (int): Number input channels. |
| - expand (int): Expansion factor used to calculate the hidden_size of LSTM. |
| - num_layers (int): Number of dual-path layers. |
| """ |
|
|
| def __init__(self, channels, expand=1, num_layers=6): |
| super(SeparationNet, self).__init__() |
|
|
| self.num_layers = num_layers |
|
|
| self.dp_modules = nn.ModuleList([ |
| DualPathRNN(channels * (2 if i % 2 == 1 else 1), expand) for i in range(num_layers) |
| ]) |
|
|
| self.feature_conversion = nn.ModuleList([ |
| FeatureConversion(channels * 2, inverse=False if i % 2 == 0 else True) for i in range(num_layers) |
| ]) |
|
|
| def forward(self, x): |
| for i in range(self.num_layers): |
| x = self.dp_modules[i](x) |
| x = self.feature_conversion[i](x) |
| return x |
|
|
