75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
def conv_block(in_channels, out_channels, kernel_size=3, dilation=1):
|
|
return nn.Sequential(
|
|
nn.Conv1d(
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size=kernel_size,
|
|
dilation=dilation,
|
|
padding=(kernel_size // 2) * dilation
|
|
),
|
|
nn.InstanceNorm1d(out_channels),
|
|
nn.PReLU()
|
|
)
|
|
|
|
class AttentionBlock(nn.Module):
|
|
"""
|
|
Simple Channel Attention Block. Learns to weight channels based on their importance.
|
|
"""
|
|
def __init__(self, channels):
|
|
super(AttentionBlock, self).__init__()
|
|
self.attention = nn.Sequential(
|
|
nn.Conv1d(channels, channels // 4, kernel_size=1),
|
|
nn.ReLU(inplace=True),
|
|
nn.Conv1d(channels // 4, channels, kernel_size=1),
|
|
nn.Sigmoid()
|
|
)
|
|
|
|
def forward(self, x):
|
|
attention_weights = self.attention(x)
|
|
return x * attention_weights
|
|
|
|
class ResidualInResidualBlock(nn.Module):
|
|
def __init__(self, channels, num_convs=3):
|
|
super(ResidualInResidualBlock, self).__init__()
|
|
|
|
self.conv_layers = nn.Sequential(
|
|
*[conv_block(channels, channels) for _ in range(num_convs)]
|
|
)
|
|
|
|
self.attention = AttentionBlock(channels)
|
|
|
|
def forward(self, x):
|
|
residual = x
|
|
x = self.conv_layers(x)
|
|
x = self.attention(x)
|
|
return x + residual
|
|
|
|
class SISUGenerator(nn.Module):
|
|
def __init__(self, channels=64, num_rirb=8, alpha=1.0):
|
|
super(SISUGenerator, self).__init__()
|
|
self.alpha = alpha
|
|
|
|
self.conv1 = nn.Sequential(
|
|
nn.Conv1d(1, channels, kernel_size=7, padding=3),
|
|
nn.InstanceNorm1d(channels),
|
|
nn.PReLU(),
|
|
)
|
|
|
|
self.rir_blocks = nn.Sequential(
|
|
*[ResidualInResidualBlock(channels) for _ in range(num_rirb)]
|
|
)
|
|
|
|
self.final_layer = nn.Conv1d(channels, 1, kernel_size=3, padding=1)
|
|
|
|
def forward(self, x):
|
|
residual_input = x
|
|
x = self.conv1(x)
|
|
x_rirb_out = self.rir_blocks(x)
|
|
learned_residual = self.final_layer(x_rirb_out)
|
|
output = residual_input + self.alpha * learned_residual
|
|
|
|
return output
|