:albemic: | Experimenting with other model layouts.

generator.py

@@ -7,30 +7,46 @@ def conv_block(in_channels, out_channels, kernel_size=3, dilation=1):
         nn.PReLU()
     )
 
+class AttentionBlock(nn.Module):
+    def __init__(self, channels):
+        super(AttentionBlock, self).__init__()
+        self.attention = nn.Sequential(
+            nn.Conv1d(channels, channels // 4, kernel_size=1),
+            nn.ReLU(),
+            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):
+    def __init__(self, layer=4, num_rirb=4): #increased base layer and rirb amounts
         super(SISUGenerator, self).__init__()
-        layer = 4 # Increased base layer count
         self.conv1 = nn.Sequential(
             nn.Conv1d(1, layer, kernel_size=7, padding=3),
             nn.BatchNorm1d(layer),
             nn.PReLU(),
         )
-        self.conv_blocks = nn.Sequential(
-            conv_block(layer, layer, kernel_size=3, dilation=1),      # Local details
-            conv_block(layer, layer*2, kernel_size=5, dilation=2),    # Local Context
-            conv_block(layer*2, layer*2, kernel_size=3, dilation=16),   # Longer range dependencies
-            conv_block(layer*2, layer*2, kernel_size=5, dilation=8),    # Wider context
-            conv_block(layer*2, layer, kernel_size=5, dilation=2),    # Local Context
-            conv_block(layer, layer, kernel_size=3, dilation=1),      # Local details
-        )
-        self.final_layer = nn.Sequential(
-            nn.Conv1d(layer, 1, kernel_size=3, padding=1),
-        )
+        self.rir_blocks = nn.Sequential(*[ResidualInResidualBlock(layer) for _ in range(num_rirb)])
+        self.final_layer = nn.Conv1d(layer, 1, kernel_size=3, padding=1)
 
     def forward(self, x):
         residual = x
         x = self.conv1(x)
-        x = self.conv_blocks(x)
+        x = self.rir_blocks(x)
         x = self.final_layer(x)
         return x + residual