import torch import torch.nn as nn import torch.nn.functional as F import models.frcrn_se.complex_nn as complex_nn from models.frcrn_se.se_layer import SELayer class Encoder(nn.Module): """ Encoder module for a neural network, responsible for downsampling input features. This module consists of a convolutional layer followed by batch normalization and a Leaky ReLU activation. Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. kernel_size (tuple): Size of the convolutional kernel. stride (tuple): Stride of the convolution. padding (tuple, optional): Padding for the convolution. If None, 'SAME' padding is applied. complex (bool, optional): If True, use complex convolution layers. Default is False. padding_mode (str, optional): Padding mode for convolution. Default is "zeros". """ def __init__(self, in_channels, out_channels, kernel_size, stride, padding=None, complex=False, padding_mode="zeros"): super().__init__() # Determine padding for 'SAME' padding if not provided if padding is None: padding = [(i - 1) // 2 for i in kernel_size] # Select convolution and batch normalization layers based on complex flag if complex: conv = complex_nn.ComplexConv2d bn = complex_nn.ComplexBatchNorm2d else: conv = nn.Conv2d bn = nn.BatchNorm2d # Define convolutional layer, batch normalization, and activation function self.conv = conv(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode=padding_mode) self.bn = bn(out_channels) self.relu = nn.LeakyReLU(inplace=True) def forward(self, x): """ Forward pass through the encoder. Args: x (torch.Tensor): Input tensor of shape (B, C, H, W) where B is batch size, C is the number of channels, H is height, and W is width. Returns: torch.Tensor: Output tensor after applying convolution, batch normalization, and activation. """ x = self.conv(x) # Apply convolution x = self.bn(x) # Apply batch normalization x = self.relu(x) # Apply Leaky ReLU activation return x class Decoder(nn.Module): """ Decoder module for a neural network, responsible for upsampling input features. This module consists of a transposed convolutional layer followed by batch normalization and a Leaky ReLU activation. Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. kernel_size (tuple): Size of the transposed convolutional kernel. stride (tuple): Stride of the transposed convolution. padding (tuple, optional): Padding for the transposed convolution. Default is (0, 0). complex (bool, optional): If True, use complex transposed convolution layers. Default is False. """ def __init__(self, in_channels, out_channels, kernel_size, stride, padding=(0, 0), complex=False): super().__init__() # Select transposed convolution and batch normalization layers based on complex flag if complex: tconv = complex_nn.ComplexConvTranspose2d bn = complex_nn.ComplexBatchNorm2d else: tconv = nn.ConvTranspose2d bn = nn.BatchNorm2d # Define transposed convolutional layer, batch normalization, and activation function self.transconv = tconv(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding) self.bn = bn(out_channels) self.relu = nn.LeakyReLU(inplace=True) def forward(self, x): """ Forward pass through the decoder. Args: x (torch.Tensor): Input tensor of shape (B, C, H, W) where B is batch size, C is the number of channels, H is height, and W is width. Returns: torch.Tensor: Output tensor after applying transposed convolution, batch normalization, and activation. """ x = self.transconv(x) # Apply transposed convolution x = self.bn(x) # Apply batch normalization x = self.relu(x) # Apply Leaky ReLU activation return x class UNet(nn.Module): """ U-Net architecture for handling both real and complex inputs. This model uses an encoder-decoder structure with skip connections between corresponding encoder and decoder layers. Squeeze-and-Excitation (SE) layers are integrated into the network for channel attention enhancement. Args: input_channels (int, optional): Number of input channels. Default is 1. complex (bool, optional): If True, use complex layers. Default is False. model_complexity (int, optional): Determines the number of channels in the model. Default is 45. model_depth (int, optional): Depth of the U-Net model (number of encoder/decoder pairs). Default is 20. padding_mode (str, optional): Padding mode for convolutions. Default is "zeros". """ def __init__(self, input_channels=1, complex=False, model_complexity=45, model_depth=20, padding_mode="zeros"): super().__init__() # Adjust model complexity for complex models if complex: model_complexity = int(model_complexity // 1.414) # Initialize model parameters based on specified complexity and depth self.set_size(model_complexity=model_complexity, input_channels=input_channels, model_depth=model_depth) self.encoders = [] self.model_length = model_depth // 2 self.fsmn = complex_nn.ComplexUniDeepFsmn(128, 128, 128) self.se_layers_enc = [] self.fsmn_enc = [] # Build the encoder structure for i in range(self.model_length): fsmn_enc = complex_nn.ComplexUniDeepFsmn_L1(128, 128, 128) self.add_module("fsmn_enc{}".format(i), fsmn_enc) self.fsmn_enc.append(fsmn_enc) module = Encoder(self.enc_channels[i], self.enc_channels[i + 1], kernel_size=self.enc_kernel_sizes[i], stride=self.enc_strides[i], padding=self.enc_paddings[i], complex=complex, padding_mode=padding_mode) self.add_module("encoder{}".format(i), module) self.encoders.append(module) se_layer_enc = SELayer(self.enc_channels[i + 1], 8) self.add_module("se_layer_enc{}".format(i), se_layer_enc) self.se_layers_enc.append(se_layer_enc) # Build the decoder structure self.decoders = [] self.fsmn_dec = [] self.se_layers_dec = [] for i in range(self.model_length): fsmn_dec = complex_nn.ComplexUniDeepFsmn_L1(128, 128, 128) self.add_module("fsmn_dec{}".format(i), fsmn_dec) self.fsmn_dec.append(fsmn_dec) module = Decoder(self.dec_channels[i] * 2, self.dec_channels[i + 1], kernel_size=self.dec_kernel_sizes[i], stride=self.dec_strides[i], padding=self.dec_paddings[i], complex=complex) self.add_module("decoder{}".format(i), module) self.decoders.append(module) if i < self.model_length - 1: se_layer_dec = SELayer(self.dec_channels[i + 1], 8) self.add_module("se_layer_dec{}".format(i), se_layer_dec) self.se_layers_dec.append(se_layer_dec) # Define final linear layer based on complex flag if complex: conv = complex_nn.ComplexConv2d else: conv = nn.Conv2d linear = conv(self.dec_channels[-1], 1, 1) # Final layer to output desired channels self.add_module("linear", linear) self.complex = complex self.padding_mode = padding_mode # Convert lists to ModuleLists for proper parameter registration self.decoders = nn.ModuleList(self.decoders) self.encoders = nn.ModuleList(self.encoders) self.se_layers_enc = nn.ModuleList(self.se_layers_enc) self.se_layers_dec = nn.ModuleList(self.se_layers_dec) self.fsmn_enc = nn.ModuleList(self.fsmn_enc) self.fsmn_dec = nn.ModuleList(self.fsmn_dec) def forward(self, inputs): """ Forward pass for the UNet model. This method processes the input tensor through the encoder-decoder architecture, applying convolutional layers, FSMNs, and SE layers. Skip connections are used to merge features from the encoder to the decoder. Args: inputs (torch.Tensor): Input tensor of shape (batch_size, channels, height, width). Returns: torch.Tensor: Output tensor after processing, representing the computed features. """ x = inputs # Initialize input tensor xs = [] # List to store input tensors for skip connections xs_se = [] # List to store outputs after applying SE layers xs_se.append(x) # Add the initial input to the SE outputs list # Forward pass through the encoder layers for i, encoder in enumerate(self.encoders): xs.append(x) # Store the current input for skip connections if i > 0: x = self.fsmn_enc[i](x) # Apply FSMN if not the first encoder x = encoder(x) # Apply the encoder layer xs_se.append(self.se_layers_enc[i](x)) # Apply SE layer and store the result x = self.fsmn(x) # Apply the final FSMN after all encoders p = x # Initialize output tensor for decoders # Forward pass through the decoder layers for i, decoder in enumerate(self.decoders): p = decoder(p) # Apply the decoder layer if i < self.model_length - 1: p = self.fsmn_dec[i](p) # Apply FSMN if not the last decoder if i == self.model_length - 1: break # Stop processing at the last decoder layer if i < self.model_length - 2: p = self.se_layers_dec[i](p) # Apply SE layer for intermediate decoders p = torch.cat([p, xs_se[self.model_length - 1 - i]], dim=1) # Concatenate skip connection # Final output processing # cmp_spec: [batch, 1, 513, 64, 2] cmp_spec = self.linear(p) # Apply linear transformation to produce final output return cmp_spec # Return the computed output tensor def set_size(self, model_complexity, model_depth=20, input_channels=1): """ Set the architecture parameters for the UNet model based on specified complexity and depth. This method configures the encoder and decoder layers of the UNet by setting the number of channels, kernel sizes, strides, and paddings for each layer according to the provided model complexity and depth. Args: model_complexity (int): Base number of channels for the model. model_depth (int, optional): Depth of the UNet model, determining the number of encoder/decoder pairs. Default is 20. input_channels (int, optional): Number of input channels to the model. Default is 1. Raises: ValueError: If an unknown model depth is provided. """ # Configuration for model depth of 14 if model_depth == 14: # Set encoder channels for model depth of 14 self.enc_channels = [input_channels, 128, 128, 128, 128, 128, 128, 128] # Define kernel sizes for encoder layers self.enc_kernel_sizes = [(5, 2), (5, 2), (5, 2), (5, 2), (5, 2), (5, 2), (2, 2)] # Define strides for encoder layers self.enc_strides = [(2, 1), (2, 1), (2, 1), (2, 1), (2, 1), (2, 1), (2, 1)] # Define paddings for encoder layers self.enc_paddings = [(0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1)] # Set decoder channels for model depth of 14 self.dec_channels = [64, 128, 128, 128, 128, 128, 128, 1] # Define kernel sizes for decoder layers self.dec_kernel_sizes = [(2, 2), (5, 2), (5, 2), (5, 2), (6, 2), (5, 2), (5, 2)] # Define strides for decoder layers self.dec_strides = [(2, 1), (2, 1), (2, 1), (2, 1), (2, 1), (2, 1), (2, 1)] # Define paddings for decoder layers self.dec_paddings = [(0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1), (0, 1)] # Configuration for model depth of 20 elif model_depth == 20: # Set encoder channels for model depth of 20 self.enc_channels = [input_channels, model_complexity, model_complexity, model_complexity * 2, model_complexity * 2, model_complexity * 2, model_complexity * 2, model_complexity * 2, model_complexity * 2, model_complexity * 2, 128] # Define kernel sizes for encoder layers self.enc_kernel_sizes = [(7, 1), (1, 7), (6, 4), (7, 5), (5, 3), (5, 3), (5, 3), (5, 3), (5, 3), (5, 3)] # Define strides for encoder layers self.enc_strides = [(1, 1), (1, 1), (2, 2), (2, 1), (2, 2), (2, 1), (2, 2), (2, 1), (2, 2), (2, 1)] # Define paddings for encoder layers self.enc_paddings = [(3, 0), (0, 3), None, # None padding for certain layers None, None, # Adjusted padding based on layer requirements None, None, None, None, None] # Set decoder channels for model depth of 20 self.dec_channels = [0, model_complexity * 2, model_complexity * 2, model_complexity * 2, model_complexity * 2, model_complexity * 2, model_complexity * 2, model_complexity * 2, model_complexity * 2, model_complexity * 2, model_complexity * 2, model_complexity * 2] # Define kernel sizes for decoder layers self.dec_kernel_sizes = [(4, 3), (4, 2), (4, 3), (4, 2), (4, 3), (4, 2), (6, 3), (7, 4), (1, 7), (7, 1)] # Define strides for decoder layers self.dec_strides = [(2, 1), (2, 2), (2, 1), (2, 2), (2, 1), (2, 2), (2, 1), (2, 2), (1, 1), (1, 1)] # Define paddings for decoder layers self.dec_paddings = [(1, 1), (1, 0), (1, 1), (1, 0), (1, 1), (1, 0), (2, 1), (2, 1), (0, 3), (3, 0)] else: # Raise an error if an unknown model depth is specified raise ValueError("Unknown model depth : {}".format(model_depth))