⚗️ | Experimenting with training optimization.

.gitignore

@@ -164,4 +164,5 @@ cython_debug/
 backup/
 dataset/
 old-output/
+output/
 *.wav

discriminator.py

@@ -1,5 +1,4 @@
 import torch.nn as nn
-import torch
 
 class SISUDiscriminator(nn.Module):
     def __init__(self):

generator.py

@@ -5,13 +5,18 @@ class SISUGenerator(nn.Module):
         super(SISUGenerator, self).__init__()
         self.model = nn.Sequential(
             nn.Conv1d(2, 128, kernel_size=3, padding=1),
+            nn.LeakyReLU(0.2, inplace=True),
             nn.Conv1d(128, 256, kernel_size=3, padding=1),
+            nn.LeakyReLU(0.2, inplace=True),
 
             nn.Upsample(scale_factor=upscale_scale, mode='nearest'),
 
             nn.Conv1d(256, 128, kernel_size=3, padding=1),
+            nn.LeakyReLU(0.2, inplace=True),
             nn.Conv1d(128, 64, kernel_size=3, padding=1),
-            nn.Conv1d(64, 2, kernel_size=3, padding=1)
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv1d(64, 2, kernel_size=3, padding=1),
+            nn.Tanh()
         )
 
     def forward(self, x):

training.py

@@ -1,6 +1,8 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
+
+import torch.nn.functional as F
 import torchaudio
 import tqdm
 
@@ -25,8 +27,8 @@ val_size = int(dataset_size-train_size)
 
 train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
 
-train_data_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)
-val_data_loader = DataLoader(val_dataset, batch_size=1, shuffle=True)
+train_data_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
+val_data_loader = DataLoader(val_dataset, batch_size=8, shuffle=True)
 
 # Initialize models and move them to device
 generator = SISUGenerator()
@@ -43,65 +45,91 @@ criterion_d = nn.BCEWithLogitsLoss()
 optimizer_g = optim.Adam(generator.parameters(), lr=0.0001, betas=(0.5, 0.999))
 optimizer_d = optim.Adam(discriminator.parameters(), lr=0.0001, betas=(0.5, 0.999))
 
-# Training loop
-num_epochs = 500
+# Scheduler
+scheduler_g = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_g, mode='min', factor=0.5, patience=5)
+scheduler_d = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_d, mode='min', factor=0.5, patience=5)
 
-for epoch in range(num_epochs):
-    low_quality_audio = torch.empty((1))
-    high_quality_audio = torch.empty((1))
-    ai_enhanced_audio = torch.empty((1))
-    total_d_loss = 0
-    total_g_loss = 0
+def snr(y_true, y_pred):
+    noise = y_true - y_pred
+    signal_power = torch.mean(y_true ** 2)
+    noise_power = torch.mean(noise ** 2)
+    snr_db = 10 * torch.log10(signal_power / noise_power)
+    return snr_db
 
-    # Training
-    for low_quality, high_quality in tqdm.tqdm(train_data_loader, desc=f"Epoch {epoch+1}/{num_epochs}"):
-        high_quality = high_quality.to(device)
-        low_quality = low_quality.to(device)
+def discriminator_train(discriminator, optimizer, criterion, generator, real_labels, fake_labels, high_quality, low_quality):
+    optimizer.zero_grad()
 
-        batch_size = 1
-        real_labels = torch.ones(batch_size, 1).to(device)
-        fake_labels = torch.zeros(batch_size, 1).to(device)
+    discriminator_decision_from_real = discriminator(high_quality)
+    d_loss_real = criterion(discriminator_decision_from_real, real_labels)
 
-        ###### Train Discriminator ######
-        discriminator.train()
-        optimizer_d.zero_grad()
+    generator_output = generator(low_quality)
+    discriminator_decision_from_fake = discriminator(generator_output.detach())
+    d_loss_fake = criterion(discriminator_decision_from_fake, fake_labels)
 
-        # 1. Real data
-        real_outputs = discriminator(high_quality)
-        d_loss_real = criterion_d(real_outputs, real_labels)
+    d_loss = (d_loss_real + d_loss_fake) / 2.0
 
-        # 2. Fake data
-        fake_audio = generator(low_quality)
-        fake_outputs = discriminator(fake_audio.detach())
-        d_loss_fake = criterion_d(fake_outputs, fake_labels)
+    d_loss.backward()
+    nn.utils.clip_grad_norm_(discriminator.parameters(), max_norm=1.0) #Gradient Clipping
+    optimizer.step()
+    # print(f"Discriminator Loss: {d_loss.item():.4f}, Mean Real Logit: {discriminator_decision_from_real.mean().item():.2f}, Mean Fake Logit: {discriminator_decision_from_fake.mean().item():.2f}")
 
-        d_loss = (d_loss_real + d_loss_fake) / 2.0  # Without gradient penalty
-        d_loss.backward()
-        optimizer_d.step()
-        total_d_loss += d_loss.item()
+def start_training():
 
-        generator.train()
-        optimizer_g.zero_grad()
+    # Training loop
+    # discriminator_epochs = 1000
+    generator_epochs = 500
+    for generator_epoch in range(generator_epochs):
+        low_quality_audio = torch.empty((1))
+        high_quality_audio = torch.empty((1))
+        ai_enhanced_audio = torch.empty((1))
 
-        # Generator loss: how well fake data fools the discriminator
-        fake_outputs = discriminator(fake_audio)  # No detach here
-        g_loss = criterion_g(fake_outputs, real_labels)  # Train generator to produce real-like outputs
+        # Training
+        for low_quality, high_quality in tqdm.tqdm(train_data_loader, desc=f"Epoch {generator_epoch+1}/{generator_epochs}"):
+            high_quality = high_quality.to(device)
+            low_quality = low_quality.to(device)
 
-        g_loss.backward()
-        optimizer_g.step()
-        total_g_loss += g_loss.item()
+            batch_size = high_quality.size(0)
+            real_labels = torch.ones(batch_size, 1).to(device)
+            fake_labels = torch.zeros(batch_size, 1).to(device)
 
-        low_quality_audio = low_quality
-        high_quality_audio = high_quality
-        ai_enhanced_audio = fake_audio
+            # Train Discriminator
+            discriminator.train()
 
-    if epoch % 10 == 0:
-        print(f"Saved epoch {epoch}!")
-        torchaudio.save(f"./output/epoch-{epoch}-audio-crap.wav", low_quality_audio[0].cpu(), 44100)
-        torchaudio.save(f"./output/epoch-{epoch}-audio-ai.wav", ai_enhanced_audio[0].cpu(), 44100)
-        torchaudio.save(f"./output/epoch-{epoch}-audio-orig.wav", high_quality_audio[0].cpu(), 44100)
+            for _ in range(3):
+                discriminator_train(discriminator, optimizer_d, criterion_d, generator, real_labels, fake_labels, high_quality, low_quality)
 
-torch.save(generator.state_dict(), "generator.pt")
-torch.save(discriminator.state_dict(), "discriminator.pt")
+            # Train Generator
+            generator.train()
+            optimizer_g.zero_grad()
 
-print("Training complete!")
+            # Generator loss: how well fake data fools the discriminator
+            generator_output = generator(low_quality)
+            discriminator_decision = discriminator(generator_output)  # No detach here
+            g_loss = criterion_g(discriminator_decision, real_labels)  # Train generator to produce real-like outputs
+
+            g_loss.backward()
+            optimizer_g.step()
+
+            low_quality_audio = low_quality
+            high_quality_audio = high_quality
+            ai_enhanced_audio = generator_output
+
+        metric = snr(high_quality_audio, ai_enhanced_audio)
+        print(f"Generator metric {metric}!")
+        scheduler_g.step(metric)
+
+        if generator_epoch % 10 == 0:
+            print(f"Saved epoch {generator_epoch}!")
+            torchaudio.save(f"./output/epoch-{generator_epoch}-audio-crap.wav", low_quality_audio[0].cpu(), 44100)
+            torchaudio.save(f"./output/epoch-{generator_epoch}-audio-ai.wav", ai_enhanced_audio[0].cpu(), 44100)
+            torchaudio.save(f"./output/epoch-{generator_epoch}-audio-orig.wav", high_quality_audio[0].cpu(), 44100)
+
+        if generator_epoch % 50 == 0:
+            torch.save(discriminator.state_dict(), "discriminator.pt")
+            torch.save(generator.state_dict(), "generator.pt")
+
+    torch.save(discriminator.state_dict(), "discriminator.pt")
+    torch.save(generator.state_dict(), "generator.pt")
+    print("Training complete!")
+
+start_training()