Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class UniDeepFsmn(nn.Module): | |
| """ | |
| A single layer Deep Feedforward Sequential Memory Network (FSMN) for unidirectional processing. | |
| This model uses a combination of linear layers and convolutional layers to process input features. | |
| Attributes: | |
| - input_dim (int): Number of input features. | |
| - output_dim (int): Number of output features. | |
| - lorder (int): Order of the linear filter. | |
| - hidden_size (int): Number of hidden units in the linear layer. | |
| """ | |
| def __init__(self, input_dim, output_dim, lorder=None, hidden_size=None): | |
| super(UniDeepFsmn, self).__init__() | |
| self.input_dim = input_dim | |
| self.output_dim = output_dim | |
| if lorder is None: | |
| return # If no filter order is provided, skip initialization | |
| self.lorder = lorder | |
| self.hidden_size = hidden_size | |
| # Linear layer to project input to hidden space | |
| self.linear = nn.Linear(input_dim, hidden_size) | |
| # Projection layer to output space | |
| self.project = nn.Linear(hidden_size, output_dim, bias=False) | |
| # Depthwise convolution layer for filtering | |
| self.conv1 = nn.Conv2d(output_dim, output_dim, [lorder, 1], [1, 1], groups=output_dim, bias=False) | |
| def forward(self, input): | |
| """ | |
| Forward pass through the UniDeepFsmn model. | |
| Parameters: | |
| - input (Tensor): Input tensor of shape (batch_size, sequence_length, input_dim). | |
| Returns: | |
| - Tensor: Output tensor after processing, with the same shape as input. | |
| """ | |
| # Apply linear transformation and ReLU activation | |
| f1 = F.relu(self.linear(input)) | |
| # Project to output dimension | |
| p1 = self.project(f1) | |
| # Reshape and pad the tensor for convolution | |
| x = torch.unsqueeze(p1, 1) # Shape: (b, c, T, h) | |
| x_per = x.permute(0, 3, 2, 1) # Permute to shape (b, h, T, c) | |
| y = F.pad(x_per, [0, 0, self.lorder - 1, 0]) # Pad the tensor | |
| # Add convolutional output to original input | |
| out = x_per + self.conv1(y) | |
| out1 = out.permute(0, 3, 2, 1) # Restore original shape | |
| return input + out1.squeeze() # Return the combined output | |
| class ComplexUniDeepFsmn(nn.Module): | |
| """ | |
| A complex variant of the UniDeepFsmn that processes complex-valued input. | |
| This model has separate layers for the real and imaginary components. | |
| Attributes: | |
| - nIn (int): Number of input features. | |
| - nHidden (int): Number of hidden units in the FSMN layers. | |
| - nOut (int): Number of output features. | |
| """ | |
| def __init__(self, nIn, nHidden=128, nOut=128): | |
| super(ComplexUniDeepFsmn, self).__init__() | |
| # Initialize FSMN layers for real and imaginary parts | |
| self.fsmn_re_L1 = UniDeepFsmn(nIn, nHidden, 20, nHidden) | |
| self.fsmn_im_L1 = UniDeepFsmn(nIn, nHidden, 20, nHidden) | |
| self.fsmn_re_L2 = UniDeepFsmn(nHidden, nOut, 20, nHidden) | |
| self.fsmn_im_L2 = UniDeepFsmn(nHidden, nOut, 20, nHidden) | |
| def forward(self, x): | |
| """ | |
| Forward pass through the ComplexUniDeepFsmn model. | |
| Parameters: | |
| - x (Tensor): Input tensor of shape (batch_size, channels, height, time, 2) where | |
| the last dimension represents the real and imaginary parts. | |
| Returns: | |
| - Tensor: Output tensor after processing. | |
| """ | |
| # Reshape input to [b, c*h, T, 2] | |
| b, c, h, T, d = x.size() | |
| x = torch.reshape(x, (b, c * h, T, d)) # Flatten channel and height | |
| x = torch.transpose(x, 1, 2) # Permute to [b, T, c*h, 2] | |
| # Process the real and imaginary parts | |
| real_L1 = self.fsmn_re_L1(x[..., 0]) - self.fsmn_im_L1(x[..., 1]) | |
| imaginary_L1 = self.fsmn_re_L1(x[..., 1]) + self.fsmn_im_L1(x[..., 0]) | |
| # Process the second layer | |
| real = self.fsmn_re_L2(real_L1) - self.fsmn_im_L2(imaginary_L1) | |
| imaginary = self.fsmn_re_L2(imaginary_L1) + self.fsmn_im_L2(real_L1) | |
| # Combine real and imaginary parts into output tensor | |
| output = torch.stack((real, imaginary), dim=-1) # Shape: [b, T, h, 2] | |
| output = torch.transpose(output, 1, 2) # Shape: [b, h, T, 2] | |
| output = torch.reshape(output, (b, c, h, T, d)) # Restore original shape | |
| return output | |
| class ComplexUniDeepFsmn_L1(nn.Module): | |
| """ | |
| A complex variant of UniDeepFsmn for the first layer. | |
| This model processes complex-valued input and has two FSMN layers for the real and imaginary parts. | |
| Attributes: | |
| - nIn (int): Number of input features. | |
| - nHidden (int): Number of hidden units in the FSMN layers. | |
| - nOut (int): Number of output features. | |
| """ | |
| def __init__(self, nIn, nHidden=128, nOut=128): | |
| super(ComplexUniDeepFsmn_L1, self).__init__() | |
| self.fsmn_re_L1 = UniDeepFsmn(nIn, nHidden, 20, nHidden) | |
| self.fsmn_im_L1 = UniDeepFsmn(nIn, nHidden, 20, nHidden) | |
| def forward(self, x): | |
| """ | |
| Forward pass through the ComplexUniDeepFsmn_L1 model. | |
| Parameters: | |
| - x (Tensor): Input tensor of shape (batch_size, channels, height, time, 2). | |
| Returns: | |
| - Tensor: Output tensor after processing. | |
| """ | |
| b, c, h, T, d = x.size() | |
| x = torch.transpose(x, 1, 3) # Shape: [b, T, h, c, 2] | |
| x = torch.reshape(x, (b * T, h, c, d)) # Reshape to process | |
| # Process the real and imaginary parts | |
| real = self.fsmn_re_L1(x[..., 0]) - self.fsmn_im_L1(x[..., 1]) | |
| imaginary = self.fsmn_re_L1(x[..., 1]) + self.fsmn_im_L1(x[..., 0]) | |
| # Combine results and reshape back to original dimensions | |
| output = torch.stack((real, imaginary), dim=-1) # Shape: [b*T, h, c, 2] | |
| output = torch.reshape(output, (b, T, h, c, d)) # Restore shape to [b, T, h, c, 2] | |
| output = torch.transpose(output, 1, 3) # Shape: [b, c, h, T, 2] | |
| return output | |
| class BidirectionalLSTM_L1(nn.Module): | |
| """ | |
| A unidirectional LSTM model for processing sequences. | |
| Attributes: | |
| - nIn (int): Number of input features. | |
| - nHidden (int): Number of hidden units in the LSTM. | |
| - nOut (int): Number of output features. | |
| """ | |
| def __init__(self, nIn, nHidden, nOut): | |
| super(BidirectionalLSTM_L1, self).__init__() | |
| self.rnn = nn.GRU(nIn, nHidden, bidirectional=False) # Using GRU instead of LSTM | |
| def forward(self, input): | |
| """ | |
| Forward pass through the BidirectionalLSTM_L1 model. | |
| Parameters: | |
| - input (Tensor): Input tensor of shape (sequence_length, batch_size, input_dim). | |
| Returns: | |
| - Tensor: Output tensor after processing. | |
| """ | |
| output, _ = self.rnn(input) # Forward pass through GRU | |
| return output | |
| class BidirectionalLSTM_L2(nn.Module): | |
| """ | |
| A unidirectional Long Short-Term Memory (LSTM) network that processes input sequences | |
| and produces an output using a linear embedding layer. | |
| Attributes: | |
| rnn (nn.GRU): The GRU layer for processing the input sequences. | |
| embedding (nn.Linear): A linear layer that transforms the output of the GRU to the desired output dimension. | |
| """ | |
| def __init__(self, nIn, nHidden, nOut): | |
| super(BidirectionalLSTM_L2, self).__init__() | |
| # Initialize the GRU layer | |
| self.rnn = nn.GRU(nIn, nHidden, bidirectional=False) | |
| # Initialize the linear embedding layer | |
| self.embedding = nn.Linear(nHidden, nOut) | |
| def forward(self, input): | |
| """ | |
| Forward pass through the Bidirectional LSTM network. | |
| Args: | |
| input (torch.Tensor): Input tensor of shape (T, b, nIn), where T is the sequence length, | |
| b is the batch size, and nIn is the input feature size. | |
| Returns: | |
| torch.Tensor: Output tensor of shape (T, b, nOut), where nOut is the output feature size. | |
| """ | |
| recurrent, _ = self.rnn(input) # Process the input through the GRU layer | |
| # Get the shape of the recurrent output | |
| T, b, h = recurrent.size() | |
| t_rec = recurrent.view(T * b, h) # Flatten the output for the embedding layer | |
| output = self.embedding(t_rec) # Apply the linear embedding layer | |
| output = output.view(T, b, -1) # Reshape back to (T, b, nOut) | |
| return output | |
| class ComplexBidirectionalLSTM(nn.Module): | |
| """ | |
| A complex-valued bidirectional LSTM that processes input sequences containing real | |
| and imaginary components, producing a complex-valued output. | |
| Attributes: | |
| lstm_re_L1 (BidirectionalLSTM_L1): First layer for the real part of the input. | |
| lstm_im_L1 (BidirectionalLSTM_L1): First layer for the imaginary part of the input. | |
| lstm_re_L2 (BidirectionalLSTM_L2): Second layer for the real part of the input. | |
| lstm_im_L2 (BidirectionalLSTM_L2): Second layer for the imaginary part of the input. | |
| """ | |
| def __init__(self, nIn, nHidden=128, nOut=1024): | |
| super(ComplexBidirectionalLSTM, self).__init__() | |
| # Initialize the first and second LSTM layers for real and imaginary components | |
| self.lstm_re_L1 = BidirectionalLSTM_L1(nIn, nHidden, nOut) | |
| self.lstm_im_L1 = BidirectionalLSTM_L1(nIn, nHidden, nOut) | |
| self.lstm_re_L2 = BidirectionalLSTM_L2(nHidden, nHidden, nOut) | |
| self.lstm_im_L2 = BidirectionalLSTM_L2(nHidden, nHidden, nOut) | |
| def forward(self, x): | |
| """ | |
| Forward pass through the complex-valued bidirectional LSTM. | |
| Args: | |
| x (torch.Tensor): Input tensor of shape (b, c, h, T, 2) where: | |
| - b is the batch size, | |
| - c is the number of channels, | |
| - h is the number of hidden units, | |
| - T is the sequence length, | |
| - 2 represents the real and imaginary parts. | |
| Returns: | |
| torch.Tensor: Output tensor of shape (b, c, h, T, 2). | |
| """ | |
| # Get the shape of the input tensor | |
| b, c, h, T, d = x.size() | |
| # Reshape the input for processing | |
| x = torch.reshape(x, (b, c*h, T, d)) | |
| # Transpose to prepare for LSTM processing | |
| x = torch.transpose(x, 0, 2) # Shape: (T, c*h, d) | |
| x = torch.transpose(x, 1, 2) # Shape: (T, d, c*h) | |
| # Process the real and imaginary parts through LSTM layers | |
| real_L1 = self.lstm_re_L1(x[..., 0]) - self.lstm_im_L1(x[..., 1]) | |
| imaginary_L1 = self.lstm_re_L1(x[..., 1]) + self.lstm_im_L1(x[..., 0]) | |
| real = self.lstm_re_L2(real_L1) - self.lstm_im_L2(imaginary_L1) | |
| imaginary = self.lstm_re_L2(imaginary_L1) + self.lstm_im_L2(real_L1) | |
| # Stack the real and imaginary parts to create the output tensor | |
| output = torch.stack((real, imaginary), dim=-1) # Shape: (T, b, h, 2) | |
| output = torch.transpose(output, 1, 2) # Shape: (T, h, b, 2) | |
| output = torch.transpose(output, 0, 2) # Shape: (b, h, T, 2) | |
| output = torch.reshape(output, (b, c, h, T, d)) # Shape: (b, c, h, T, 2) | |
| return output | |
| class ComplexConv2d(nn.Module): | |
| """ | |
| A complex-valued 2D convolutional layer that processes input tensors with real | |
| and imaginary parts, returning a complex output. | |
| Attributes: | |
| conv_re (nn.Conv2d): Convolutional layer for the real part. | |
| conv_im (nn.Conv2d): Convolutional layer for the imaginary part. | |
| """ | |
| def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, **kwargs): | |
| super().__init__() | |
| # Initialize convolutional layers for real and imaginary components | |
| self.conv_re = nn.Conv2d(in_channel, out_channel, kernel_size, stride=stride, padding=padding, | |
| dilation=dilation, groups=groups, bias=bias, **kwargs) | |
| self.conv_im = nn.Conv2d(in_channel, out_channel, kernel_size, stride=stride, padding=padding, | |
| dilation=dilation, groups=groups, bias=bias, **kwargs) | |
| def forward(self, x): | |
| """ | |
| Forward pass through the complex convolutional layer. | |
| Args: | |
| x (torch.Tensor): Input tensor of shape (batch, channel, axis1, axis2, 2) | |
| representing the real and imaginary parts. | |
| Returns: | |
| torch.Tensor: Output tensor containing the convolved real and imaginary parts. | |
| """ | |
| # Apply convolution to the real and imaginary parts | |
| real = self.conv_re(x[..., 0]) - self.conv_im(x[..., 1]) | |
| imaginary = self.conv_re(x[..., 1]) + self.conv_im(x[..., 0]) | |
| output = torch.stack((real, imaginary), dim=-1) # Stack real and imaginary components | |
| return output | |
| class ComplexConvTranspose2d(nn.Module): | |
| """ | |
| A complex-valued 2D transposed convolutional layer that processes input tensors | |
| with real and imaginary parts, returning a complex output. | |
| Attributes: | |
| tconv_re (nn.ConvTranspose2d): Transposed convolutional layer for the real part. | |
| tconv_im (nn.ConvTranspose2d): Transposed convolutional layer for the imaginary part. | |
| """ | |
| def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, output_padding=0, dilation=1, groups=1, bias=True, **kwargs): | |
| super().__init__() | |
| # Initialize transposed convolutional layers for real and imaginary components | |
| self.tconv_re = nn.ConvTranspose2d(in_channel, out_channel, | |
| kernel_size=kernel_size, | |
| stride=stride, | |
| padding=padding, | |
| output_padding=output_padding, | |
| groups=groups, | |
| bias=bias, | |
| dilation=dilation, | |
| **kwargs) | |
| self.tconv_im = nn.ConvTranspose2d(in_channel, out_channel, | |
| kernel_size=kernel_size, | |
| stride=stride, | |
| padding=padding, | |
| output_padding=output_padding, | |
| groups=groups, | |
| bias=bias, | |
| dilation=dilation, | |
| **kwargs) | |
| def forward(self, x): | |
| """ | |
| Forward pass through the complex transposed convolutional layer. | |
| Args: | |
| x (torch.Tensor): Input tensor of shape (batch, channel, axis1, axis2, 2) | |
| representing the real and imaginary parts. | |
| Returns: | |
| torch.Tensor: Output tensor containing the transposed convoluted real and imaginary parts. | |
| """ | |
| # Apply transposed convolution to the real and imaginary parts | |
| real = self.tconv_re(x[..., 0]) - self.tconv_im(x[..., 1]) | |
| imaginary = self.tconv_re(x[..., 1]) + self.tconv_im(x[..., 0]) | |
| output = torch.stack((real, imaginary), dim=-1) # Stack real and imaginary components | |
| return output | |
| class ComplexBatchNorm2d(nn.Module): | |
| """ | |
| A complex-valued batch normalization layer that normalizes input tensors with | |
| separate real and imaginary components. | |
| This layer applies batch normalization independently to the real and imaginary parts of the input, | |
| ensuring that each part is normalized appropriately. It is particularly useful in complex-valued networks, | |
| where inputs are represented as pairs of real and imaginary components. | |
| Attributes: | |
| bn_re (nn.BatchNorm2d): Batch normalization layer for the real part of the input. | |
| bn_im (nn.BatchNorm2d): Batch normalization layer for the imaginary part of the input. | |
| """ | |
| def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, | |
| track_running_stats=True, **kwargs): | |
| """ | |
| Initializes the ComplexBatchNorm2d layer. | |
| Args: | |
| num_features (int): Number of features (channels) for the input. | |
| eps (float, optional): A value added to the denominator for numerical stability. Default is 1e-5. | |
| momentum (float, optional): Momentum for the running mean and variance. Default is 0.1. | |
| affine (bool, optional): If True, this layer has learnable parameters. Default is True. | |
| track_running_stats (bool, optional): If True, track the running mean and variance. Default is True. | |
| """ | |
| super().__init__() | |
| # Initialize batch normalization layers for real and imaginary parts | |
| self.bn_re = nn.BatchNorm2d(num_features=num_features, momentum=momentum, affine=affine, eps=eps, track_running_stats=track_running_stats, **kwargs) | |
| self.bn_im = nn.BatchNorm2d(num_features=num_features, momentum=momentum, affine=affine, eps=eps, track_running_stats=track_running_stats, **kwargs) | |
| def forward(self, x): | |
| """ | |
| Forward pass through the complex batch normalization layer. | |
| Args: | |
| x (torch.Tensor): Input tensor of shape (batch, channel, height, width, 2), | |
| where the last dimension represents the real and imaginary parts. | |
| Returns: | |
| torch.Tensor: Output tensor containing the normalized real and imaginary components, | |
| with the same shape as the input tensor. | |
| """ | |
| # Apply batch normalization to the real part | |
| real = self.bn_re(x[..., 0]) | |
| # Apply batch normalization to the imaginary part | |
| imag = self.bn_im(x[..., 1]) | |
| # Stack the normalized real and imaginary parts back together | |
| output = torch.stack((real, imag), dim=-1) | |
| return output | |