From fb7b624c877d6d7d9016b75f04c04e3b39c69761 Mon Sep 17 00:00:00 2001
From: NikkeDoy <niklas.siltala@gmail.com>
Date: Mon, 10 Feb 2025 12:44:42 +0200
Subject: [PATCH] :alembic: | Experimenting with very small model.

---
 discriminator.py | 9 ++++-----
 generator.py     | 8 ++++----
 training.py      | 2 +-
 3 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/discriminator.py b/discriminator.py
index b800eda..b1d82e1 100644
--- a/discriminator.py
+++ b/discriminator.py
@@ -13,7 +13,7 @@ def discriminator_block(in_channels, out_channels, kernel_size=3, stride=1, dila
 class SISUDiscriminator(nn.Module):
     def __init__(self):
         super(SISUDiscriminator, self).__init__()
-        layers = 32 # Increased base layer count
+        layers = 4 # Increased base layer count
         self.model = nn.Sequential(
             # Initial Convolution
             discriminator_block(1, layers, kernel_size=7, stride=2, dilation=1), # Downsample
@@ -21,10 +21,9 @@ class SISUDiscriminator(nn.Module):
             # Core Discriminator Blocks with varied kernels and dilations
             discriminator_block(layers, layers * 2, kernel_size=5, stride=2, dilation=1), # Downsample
             discriminator_block(layers * 2, layers * 4, kernel_size=5, dilation=4),
-            discriminator_block(layers * 4, layers * 8, kernel_size=5, dilation=16),
-            discriminator_block(layers * 8, layers * 4, kernel_size=3, dilation=8),
-            discriminator_block(layers * 4, layers * 2, kernel_size=3, dilation=2),
-            discriminator_block(layers * 2, layers, kernel_size=5, dilation=1),
+            discriminator_block(layers * 4, layers * 4, kernel_size=5, dilation=16),
+            discriminator_block(layers * 4, layers * 2, kernel_size=3, dilation=8),
+            discriminator_block(layers * 2, layers, kernel_size=3, dilation=1),
             # Final Convolution
             discriminator_block(layers, 1, kernel_size=3, stride=1),
         )
diff --git a/generator.py b/generator.py
index 2446275..03fa279 100644
--- a/generator.py
+++ b/generator.py
@@ -10,7 +10,7 @@ def conv_block(in_channels, out_channels, kernel_size=3, dilation=1):
 class SISUGenerator(nn.Module):
     def __init__(self):
         super(SISUGenerator, self).__init__()
-        layer = 32 # Increased base layer count
+        layer = 4 # Increased base layer count
         self.conv1 = nn.Sequential(
             nn.Conv1d(1, layer, kernel_size=7, padding=3),
             nn.BatchNorm1d(layer),
@@ -18,9 +18,9 @@ class SISUGenerator(nn.Module):
         )
         self.conv_blocks = nn.Sequential(
             conv_block(layer, layer, kernel_size=3, dilation=1),      # Local details
-            conv_block(layer, layer*4, kernel_size=5, dilation=2),    # Local Context
-            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, 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
         )
diff --git a/training.py b/training.py
index e114817..6ee7116 100644
--- a/training.py
+++ b/training.py
@@ -85,7 +85,7 @@ dataset = AudioDataset(dataset_dir)
 
 # ========= SINGLE =========
 
-train_data_loader = DataLoader(dataset, batch_size=1, shuffle=True)
+train_data_loader = DataLoader(dataset, batch_size=16, shuffle=True)
 
 # Initialize models and move them to device
 generator = SISUGenerator()