⚗️ | Experimenting with larger model architecture.

generator.py

@@ -1,31 +1,39 @@
 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.BatchNorm1d(out_channels),
+        nn.PReLU()
+    )
+
 class SISUGenerator(nn.Module):
     def __init__(self):
         super(SISUGenerator, self).__init__()
-        layer = 16
-        # Convolution layers with BatchNorm and Residuals
+        layer = 32 # Increased base layer count
         self.conv1 = nn.Sequential(
-            nn.Conv1d(1, layer * 2, kernel_size=7, padding=3),
-            nn.BatchNorm1d(layer * 2),
-            nn.PReLU(),
-            nn.Conv1d(layer * 2, layer * 5, kernel_size=7, padding=3),
-            nn.BatchNorm1d(layer * 5),
-            nn.PReLU(),
-            nn.Conv1d(layer * 5, layer * 5, kernel_size=7, padding=3),
-            nn.BatchNorm1d(layer * 5),
+            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=4),    # Wider context
+            conv_block(layer*2, layer*4, kernel_size=7, dilation=8),    # Longer range dependencies
+            conv_block(layer*4, layer*4, kernel_size=3, dilation=16),   # Longer range dependencies
+            conv_block(layer*4, layer*2, kernel_size=5, dilation=8),    # Wider context
+            conv_block(layer*2, layer*2, kernel_size=3, dilation=4),    # 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 * 5, layer * 2, kernel_size=5, padding=2),
-            nn.BatchNorm1d(layer * 2),
-            nn.PReLU(),
-            nn.Conv1d(layer * 2, 1, kernel_size=3, padding=1),
-            # nn.Tanh()  # Normalize audio... if needed...
+            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.final_layer(x)
         return x + residual