215 lines
10 KiB
Python
215 lines
10 KiB
Python
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import torch as th
|
|
from torch.nn.parameter import Parameter
|
|
import numpy as np
|
|
import os
|
|
|
|
class UniDeepFsmn(nn.Module):
|
|
"""
|
|
UniDeepFsmn is a neural network module that implements a single-deep feedforward sequence memory network (FSMN).
|
|
|
|
Attributes:
|
|
input_dim (int): Dimension of the input features.
|
|
output_dim (int): Dimension of the output features.
|
|
lorder (int): Length of the order for the convolution layers.
|
|
hidden_size (int): Number of hidden units in the linear layer.
|
|
linear (nn.Linear): Linear layer to project input features to hidden size.
|
|
project (nn.Linear): Linear layer to project hidden features to output dimensions.
|
|
conv1 (nn.Conv2d): Convolutional layer for processing the output in a grouped manner.
|
|
"""
|
|
|
|
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
|
|
self.lorder = lorder
|
|
self.hidden_size = hidden_size
|
|
|
|
# Initialize the layers
|
|
self.linear = nn.Linear(input_dim, hidden_size) # Linear transformation to hidden size
|
|
self.project = nn.Linear(hidden_size, output_dim, bias=False) # Project hidden size to output dimension
|
|
self.conv1 = nn.Conv2d(output_dim, output_dim, [lorder + lorder - 1, 1], [1, 1], groups=output_dim, bias=False) # Convolution layer
|
|
|
|
def forward(self, input):
|
|
"""
|
|
Forward pass for the UniDeepFsmn model.
|
|
|
|
Args:
|
|
input (torch.Tensor): Input tensor of shape (batch_size, input_dim).
|
|
|
|
Returns:
|
|
torch.Tensor: The output tensor of the same shape as input, enhanced by the network.
|
|
"""
|
|
f1 = F.relu(self.linear(input)) # Apply linear layer followed by ReLU activation
|
|
p1 = self.project(f1) # Project to output dimension
|
|
x = th.unsqueeze(p1, 1) # Add a dimension for compatibility with Conv2d
|
|
x_per = x.permute(0, 3, 2, 1) # Permute dimensions for convolution
|
|
y = F.pad(x_per, [0, 0, self.lorder - 1, self.lorder - 1]) # Pad for causal convolution
|
|
out = x_per + self.conv1(y) # Add original input to convolution output
|
|
out1 = out.permute(0, 3, 2, 1) # Permute back to original dimensions
|
|
return input + out1.squeeze() # Return enhanced input
|
|
|
|
|
|
class UniDeepFsmn_dual(nn.Module):
|
|
"""
|
|
UniDeepFsmn_dual is a neural network module that implements a dual-deep feedforward sequence memory network (FSMN).
|
|
|
|
This class extends the UniDeepFsmn by adding a second convolution layer for richer feature extraction.
|
|
|
|
Attributes:
|
|
input_dim (int): Dimension of the input features.
|
|
output_dim (int): Dimension of the output features.
|
|
lorder (int): Length of the order for the convolution layers.
|
|
hidden_size (int): Number of hidden units in the linear layer.
|
|
linear (nn.Linear): Linear layer to project input features to hidden size.
|
|
project (nn.Linear): Linear layer to project hidden features to output dimensions.
|
|
conv1 (nn.Conv2d): First convolutional layer for processing the output.
|
|
conv2 (nn.Conv2d): Second convolutional layer for further processing the features.
|
|
"""
|
|
|
|
def __init__(self, input_dim, output_dim, lorder=None, hidden_size=None):
|
|
super(UniDeepFsmn_dual, self).__init__()
|
|
|
|
self.input_dim = input_dim
|
|
self.output_dim = output_dim
|
|
if lorder is None:
|
|
return
|
|
self.lorder = lorder
|
|
self.hidden_size = hidden_size
|
|
|
|
# Initialize the layers
|
|
self.linear = nn.Linear(input_dim, hidden_size) # Linear transformation to hidden size
|
|
self.project = nn.Linear(hidden_size, output_dim, bias=False) # Project hidden size to output dimension
|
|
self.conv1 = nn.Conv2d(output_dim, output_dim, [lorder + lorder - 1, 1], [1, 1], groups=output_dim, bias=False) # First convolution layer
|
|
self.conv2 = nn.Conv2d(output_dim, output_dim, [lorder + lorder - 1, 1], [1, 1], groups=output_dim // 4, bias=False) # Second convolution layer
|
|
|
|
def forward(self, input):
|
|
"""
|
|
Forward pass for the UniDeepFsmn_dual model.
|
|
|
|
Args:
|
|
input (torch.Tensor): Input tensor of shape (batch_size, input_dim).
|
|
|
|
Returns:
|
|
torch.Tensor: The output tensor of the same shape as input, enhanced by the network.
|
|
"""
|
|
f1 = F.relu(self.linear(input)) # Apply linear layer followed by ReLU activation
|
|
p1 = self.project(f1) # Project to output dimension
|
|
x = th.unsqueeze(p1, 1) # Add a dimension for compatibility with Conv2d
|
|
x_per = x.permute(0, 3, 2, 1) # Permute dimensions for convolution
|
|
y = F.pad(x_per, [0, 0, self.lorder - 1, self.lorder - 1]) # Pad for causal convolution
|
|
conv1_out = x_per + self.conv1(y) # Add original input to first convolution output
|
|
z = F.pad(conv1_out, [0, 0, self.lorder - 1, self.lorder - 1]) # Pad for second convolution
|
|
out = conv1_out + self.conv2(z) # Add output of second convolution
|
|
out1 = out.permute(0, 3, 2, 1) # Permute back to original dimensions
|
|
return input + out1.squeeze() # Return enhanced input
|
|
|
|
|
|
class DilatedDenseNet(nn.Module):
|
|
"""
|
|
DilatedDenseNet implements a dense network structure with dilated convolutions.
|
|
|
|
This architecture enables wider receptive fields while maintaining a lower number of parameters.
|
|
It consists of multiple convolutional layers with dilation rates that increase at each layer.
|
|
|
|
Attributes:
|
|
depth (int): Number of convolutional layers in the network.
|
|
in_channels (int): Number of input channels for the first layer.
|
|
pad (nn.ConstantPad2d): Padding layer to maintain dimensions.
|
|
twidth (int): Width of the kernel used in convolution.
|
|
kernel_size (tuple): Kernel size for convolution operations.
|
|
"""
|
|
|
|
def __init__(self, depth=4, lorder=20, in_channels=64):
|
|
super(DilatedDenseNet, self).__init__()
|
|
self.depth = depth
|
|
self.in_channels = in_channels
|
|
self.pad = nn.ConstantPad2d((1, 1, 1, 0), value=0.) # Padding for the input
|
|
self.twidth = lorder * 2 - 1 # Width of the kernel
|
|
self.kernel_size = (self.twidth, 1) # Kernel size for convolutions
|
|
|
|
# Initialize layers dynamically based on depth
|
|
for i in range(self.depth):
|
|
dil = 2 ** i # Calculate dilation rate
|
|
pad_length = lorder + (dil - 1) * (lorder - 1) - 1 # Calculate padding length
|
|
setattr(self, 'pad{}'.format(i + 1), nn.ConstantPad2d((0, 0, pad_length, pad_length), value=0.)) # Padding for dilation
|
|
setattr(self, 'conv{}'.format(i + 1),
|
|
nn.Conv2d(self.in_channels * (i + 1), self.in_channels, kernel_size=self.kernel_size,
|
|
dilation=(dil, 1), groups=self.in_channels, bias=False)) # Convolution layer with dilation
|
|
setattr(self, 'norm{}'.format(i + 1), nn.InstanceNorm2d(in_channels, affine=True)) # Normalization layer
|
|
setattr(self, 'prelu{}'.format(i + 1), nn.PReLU(self.in_channels)) # Activation layer
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass for the DilatedDenseNet model.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width).
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor after applying dense layers.
|
|
"""
|
|
skip = x # Initialize skip connection
|
|
for i in range(self.depth):
|
|
out = getattr(self, 'pad{}'.format(i + 1))(skip) # Apply padding
|
|
out = getattr(self, 'conv{}'.format(i + 1))(out) # Apply convolution
|
|
out = getattr(self, 'norm{}'.format(i + 1))(out) # Apply normalization
|
|
out = getattr(self, 'prelu{}'.format(i + 1))(out) # Apply PReLU activation
|
|
skip = th.cat([out, skip], dim=1) # Concatenate the output with the skip connection
|
|
return out # Return the final output
|
|
|
|
class UniDeepFsmn_dilated(nn.Module):
|
|
"""
|
|
UniDeepFsmn_dilated combines the UniDeepFsmn architecture with a dilated dense network
|
|
to enhance feature extraction while maintaining efficient computation.
|
|
|
|
Attributes:
|
|
input_dim (int): Dimension of the input features.
|
|
output_dim (int): Dimension of the output features.
|
|
depth (int): Depth of the dilated dense network.
|
|
lorder (int): Length of the order for the convolution layers.
|
|
hidden_size (int): Number of hidden units in the linear layer.
|
|
linear (nn.Linear): Linear layer to project input features to hidden size.
|
|
project (nn.Linear): Linear layer to project hidden features to output dimensions.
|
|
conv (DilatedDenseNet): Instance of the DilatedDenseNet for feature extraction.
|
|
"""
|
|
|
|
def __init__(self, input_dim, output_dim, lorder=None, hidden_size=None, depth=2):
|
|
super(UniDeepFsmn_dilated, self).__init__()
|
|
|
|
self.input_dim = input_dim
|
|
self.output_dim = output_dim
|
|
self.depth = depth
|
|
if lorder is None:
|
|
return
|
|
self.lorder = lorder
|
|
self.hidden_size = hidden_size
|
|
|
|
# Initialize layers
|
|
self.linear = nn.Linear(input_dim, hidden_size) # Linear transformation to hidden size
|
|
self.project = nn.Linear(hidden_size, output_dim, bias=False) # Project hidden size to output dimension
|
|
self.conv = DilatedDenseNet(depth=self.depth, lorder=lorder, in_channels=output_dim) # Dilated dense network for feature extraction
|
|
|
|
def forward(self, input):
|
|
"""
|
|
Forward pass for the UniDeepFsmn_dilated model.
|
|
|
|
Args:
|
|
input (torch.Tensor): Input tensor of shape (batch_size, input_dim).
|
|
|
|
Returns:
|
|
torch.Tensor: The output tensor of the same shape as input, enhanced by the network.
|
|
"""
|
|
f1 = F.relu(self.linear(input)) # Apply linear layer followed by ReLU activation
|
|
p1 = self.project(f1) # Project to output dimension
|
|
x = th.unsqueeze(p1, 1) # Add a dimension for compatibility with Conv2d
|
|
x_per = x.permute(0, 3, 2, 1) # Permute dimensions for convolution
|
|
out = self.conv(x_per) # Pass through the dilated dense network
|
|
out1 = out.permute(0, 3, 2, 1) # Permute back to original dimensions
|
|
|
|
return input + out1.squeeze() # Return enhanced input
|