424 lines
18 KiB
Python
424 lines
18 KiB
Python
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
|