NikkeDoy/SISU
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

⚗️ | Added some stupid ways for training + some makeup

Browse Source
This commit is contained in:
NikkeDoy
2025-10-04 22:38:11 +03:00
parent 0bc8fc2792
commit 3f23242d6f
12 changed files with 304 additions and 463 deletions
Download Patch File Download Diff File Expand all files Collapse all files

246
training.py
View File

@@ -4,25 +4,20 @@ import os
import torch
import torch.nn as nn
import torch.optim as optim
import torchaudio.transforms as T
import tqdm
from torch.amp import GradScaler, autocast
from torch.utils.data import DataLoader
from accelerate import Accelerator
from torch.utils.data import DataLoader, DistributedSampler
import training_utils
from data import AudioDataset
from discriminator import SISUDiscriminator
from generator import SISUGenerator
from training_utils import discriminator_train, generator_train
from utils.TrainingTools import discriminator_train, generator_train
# ---------------------------
# Argument parsing
# ---------------------------
parser = argparse.ArgumentParser(description="Training script (safer defaults)")
parser.add_argument("--resume", action="store_true", help="Resume training")
parser.add_argument(
"--device", type=str, default="cuda", help="Device (cuda, cpu, mps)"
)
parser.add_argument(
"--epochs", type=int, default=5000, help="Number of training epochs"
)
@@ -35,86 +30,54 @@ parser.add_argument(
args = parser.parse_args()
# ---------------------------
# Device setup
# Init accelerator
# ---------------------------
# Use requested device only if available
device = torch.device(
args.device if (args.device != "cuda" or torch.cuda.is_available()) else "cpu"
)
print(f"Using device: {device}")
# sensible performance flags
if device.type == "cuda":
torch.backends.cudnn.benchmark = True
# optional: torch.set_float32_matmul_precision("high")
debug = args.debug
# ---------------------------
# Audio transforms
# ---------------------------
sample_rate = 44100
n_fft = 1024
win_length = n_fft
hop_length = n_fft // 4
n_mels = 96
# n_mfcc = 13
# mfcc_transform = T.MFCC(
# sample_rate=sample_rate,
# n_mfcc=n_mfcc,
# melkwargs=dict(
# n_fft=n_fft,
# hop_length=hop_length,
# win_length=win_length,
# n_mels=n_mels,
# power=1.0,
# ),
# ).to(device)
mel_transform = T.MelSpectrogram(
sample_rate=sample_rate,
n_fft=n_fft,
hop_length=hop_length,
win_length=win_length,
n_mels=n_mels,
power=1.0,
).to(device)
stft_transform = T.Spectrogram(
n_fft=n_fft, win_length=win_length, hop_length=hop_length
).to(device)
# training_utils.init(mel_transform, stft_transform, mfcc_transform)
training_utils.init(mel_transform, stft_transform)
# ---------------------------
# Dataset / DataLoader
# ---------------------------
dataset_dir = "./dataset/good"
dataset = AudioDataset(dataset_dir)
train_loader = DataLoader(
dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
persistent_workers=True,
)
accelerator = Accelerator(mixed_precision="bf16")
# ---------------------------
# Models
# ---------------------------
generator = SISUGenerator().to(device)
discriminator = SISUDiscriminator().to(device)
generator = SISUGenerator()
discriminator = SISUDiscriminator()
accelerator.print("🔨 | Compiling models...")
generator = torch.compile(generator)
discriminator = torch.compile(discriminator)
accelerator.print("✅ | Compiling done!")
# ---------------------------
# Dataset / DataLoader
# ---------------------------
accelerator.print("📊 | Fetching dataset...")
dataset = AudioDataset("./dataset")
sampler = DistributedSampler(dataset) if accelerator.num_processes > 1 else None
pin_memory = torch.cuda.is_available() and not args.no_pin_memory
train_loader = DataLoader(
dataset,
sampler=sampler,
batch_size=args.batch_size,
shuffle=(sampler is None),
num_workers=args.num_workers,
pin_memory=pin_memory,
persistent_workers=pin_memory,
)
if not train_loader or not train_loader.batch_size or train_loader.batch_size == 0:
accelerator.print("🪹 | There is no data to train with! Exiting...")
exit()
loader_batch_size = train_loader.batch_size
accelerator.print("✅ | Dataset fetched!")
# ---------------------------
# Losses / Optimizers / Scalers
# ---------------------------
criterion_g = nn.BCEWithLogitsLoss()
criterion_d = nn.BCEWithLogitsLoss()
optimizer_g = optim.AdamW(
generator.parameters(), lr=0.0003, betas=(0.5, 0.999), weight_decay=0.0001
@@ -123,9 +86,6 @@ optimizer_d = optim.AdamW(
discriminator.parameters(), lr=0.0003, betas=(0.5, 0.999), weight_decay=0.0001
)
# Use modern GradScaler signature; choose device_type based on runtime device.
scaler = GradScaler(device=device)
scheduler_g = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer_g, mode="min", factor=0.5, patience=5
)
@@ -133,6 +93,17 @@ scheduler_d = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer_d, mode="min", factor=0.5, patience=5
)
criterion_g = nn.BCEWithLogitsLoss()
criterion_d = nn.MSELoss()
# ---------------------------
# Prepare accelerator
# ---------------------------
generator, discriminator, optimizer_g, optimizer_d, train_loader = accelerator.prepare(
generator, discriminator, optimizer_g, optimizer_d, train_loader
)
# ---------------------------
# Checkpoint helpers
# ---------------------------
@@ -141,44 +112,45 @@ os.makedirs(models_dir, exist_ok=True)
def save_ckpt(path, epoch):
torch.save(
{
"epoch": epoch,
"G": generator.state_dict(),
"D": discriminator.state_dict(),
"optG": optimizer_g.state_dict(),
"optD": optimizer_d.state_dict(),
"scaler": scaler.state_dict(),
"schedG": scheduler_g.state_dict(),
"schedD": scheduler_d.state_dict(),
},
path,
)
accelerator.wait_for_everyone()
if accelerator.is_main_process:
accelerator.save(
{
"epoch": epoch,
"G": accelerator.unwrap_model(generator).state_dict(),
"D": accelerator.unwrap_model(discriminator).state_dict(),
"optG": optimizer_g.state_dict(),
"optD": optimizer_d.state_dict(),
"schedG": scheduler_g.state_dict(),
"schedD": scheduler_d.state_dict(),
},
path,
)
start_epoch = 0
if args.resume:
ckpt = torch.load(os.path.join(models_dir, "last.pt"), map_location=device)
generator.load_state_dict(ckpt["G"])
discriminator.load_state_dict(ckpt["D"])
ckpt_path = os.path.join(models_dir, "last.pt")
ckpt = torch.load(ckpt_path)
accelerator.unwrap_model(generator).load_state_dict(ckpt["G"])
accelerator.unwrap_model(discriminator).load_state_dict(ckpt["D"])
optimizer_g.load_state_dict(ckpt["optG"])
optimizer_d.load_state_dict(ckpt["optD"])
scaler.load_state_dict(ckpt["scaler"])
scheduler_g.load_state_dict(ckpt["schedG"])
scheduler_d.load_state_dict(ckpt["schedD"])
start_epoch = ckpt.get("epoch", 1)
accelerator.print(f"🔁 | Resumed from epoch {start_epoch}!")
# ---------------------------
# Training loop (safer)
# ---------------------------
real_buf = torch.full(
(loader_batch_size, 1), 1, device=accelerator.device, dtype=torch.float32
)
fake_buf = torch.zeros(
(loader_batch_size, 1), device=accelerator.device, dtype=torch.float32
)
if not train_loader or not train_loader.batch_size:
print("There is no data to train with! Exiting...")
exit()
max_batch = max(1, train_loader.batch_size)
real_buf = torch.full((max_batch, 1), 0.9, device=device) # label smoothing
fake_buf = torch.zeros(max_batch, 1, device=device)
accelerator.print("🏋️ | Started training...")
try:
for epoch in range(start_epoch, args.epochs):
@@ -193,15 +165,12 @@ try:
) in enumerate(tqdm.tqdm(train_loader, desc=f"Epoch {epoch}")):
batch_size = high_quality.size(0)
high_quality = high_quality.to(device, non_blocking=True)
low_quality = low_quality.to(device, non_blocking=True)
real_labels = real_buf[:batch_size]
fake_labels = fake_buf[:batch_size]
real_labels = real_buf[:batch_size].to(accelerator.device)
fake_labels = fake_buf[:batch_size].to(accelerator.device)
# --- Discriminator ---
optimizer_d.zero_grad(set_to_none=True)
with autocast(device_type=device.type):
with accelerator.autocast():
d_loss = discriminator_train(
high_quality,
low_quality,
@@ -212,15 +181,14 @@ try:
criterion_d,
)
scaler.scale(d_loss).backward()
scaler.unscale_(optimizer_d)
torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 1.0)
scaler.step(optimizer_d)
accelerator.backward(d_loss)
torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 1)
optimizer_d.step()
# --- Generator ---
optimizer_g.zero_grad(set_to_none=True)
with autocast(device_type=device.type):
g_out, g_total, g_adv = generator_train(
with accelerator.autocast():
g_total, g_adv = generator_train(
low_quality,
high_quality,
real_labels,
@@ -229,20 +197,32 @@ try:
criterion_d,
)
scaler.scale(g_total).backward()
scaler.unscale_(optimizer_g)
torch.nn.utils.clip_grad_norm_(generator.parameters(), 1.0)
scaler.step(optimizer_g)
accelerator.backward(g_total)
torch.nn.utils.clip_grad_norm_(generator.parameters(), 1)
optimizer_g.step()
scaler.update()
d_val = accelerator.gather(d_loss.detach()).mean()
g_val = accelerator.gather(g_total.detach()).mean()
if torch.isfinite(d_val):
running_d += d_val.item()
else:
accelerator.print(
f"🫥 | NaN in discriminator loss at step {i}, skipping update."
)
if torch.isfinite(g_val):
running_g += g_val.item()
else:
accelerator.print(
f"🫥 | NaN in generator loss at step {i}, skipping update."
)
running_d += float(d_loss.detach().cpu().item())
running_g += float(g_total.detach().cpu().item())
steps += 1
# epoch averages & schedulers
if steps == 0:
print("No steps in epoch (empty dataloader?). Exiting.")
accelerator.print("🪹 | No steps in epoch (empty dataloader?). Exiting.")
break
mean_d = running_d / steps
@@ -252,22 +232,14 @@ try:
scheduler_g.step(mean_g)
save_ckpt(os.path.join(models_dir, "last.pt"), epoch)
print(f"Epoch {epoch} done | D {mean_d:.4f} | G {mean_g:.4f}")
accelerator.print(f"🤝 | Epoch {epoch} done | D {mean_d:.4f} | G {mean_g:.4f}")
except Exception:
try:
save_ckpt(os.path.join(models_dir, "crash_last.pt"), epoch)
print(f"Saved crash checkpoint for epoch {epoch}")
accelerator.print(f"💾 | Saved crash checkpoint for epoch {epoch}")
except Exception as e:
print("Failed saving crash checkpoint:", e)
accelerator.print("😬 | Failed saving crash checkpoint:", e)
raise
try:
torch.save(generator.state_dict(), os.path.join(models_dir, "final_generator.pt"))
torch.save(
discriminator.state_dict(), os.path.join(models_dir, "final_discriminator.pt")
)
except Exception as e:
print("Failed to save final states:", e)
print("Training finished.")
accelerator.print("🏁 | Training finished.")