From b7d7e95c891cbdd02f2b51204888fc261fc152fc Mon Sep 17 00:00:00 2001 From: NikkeDoy Date: Sat, 21 Dec 2024 00:24:00 +0200 Subject: [PATCH] :alembic: | Experimenting with training optimization. --- .gitignore | 1 + discriminator.py | 1 - generator.py | 7 ++- training.py | 128 +++++++++++++++++++++++++++++------------------ 4 files changed, 85 insertions(+), 52 deletions(-) diff --git a/.gitignore b/.gitignore index 2a8b716..c97a892 100644 --- a/.gitignore +++ b/.gitignore @@ -164,4 +164,5 @@ cython_debug/ backup/ dataset/ old-output/ +output/ *.wav diff --git a/discriminator.py b/discriminator.py index 9fd9e30..86a35d6 100644 --- a/discriminator.py +++ b/discriminator.py @@ -1,5 +1,4 @@ import torch.nn as nn -import torch class SISUDiscriminator(nn.Module): def __init__(self): diff --git a/generator.py b/generator.py index 8757aaf..978f02a 100644 --- a/generator.py +++ b/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): diff --git a/training.py b/training.py index 7d54123..73751f0 100644 --- a/training.py +++ b/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()