train.py

import torch
from torch import optim, nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import LambdaLR
from tqdm import tqdm
import matplotlib.pyplot as plt
import os

from data import get_splits, NSynthDataset
from miniset import load_splits
from model import VAE

# [train, validate, test] = get_splits(splits=['valid', 'test', 'train_subset']) # *** UNCOMMENT THIS LINE TO USE THE FULL DATASET ***
[train, validate] = load_splits() # This loads the mini data splits, comment this if you wish to load the full datasets using line above
test = None

for key,val in train[0].items():
    try:
        print(f"{key}: {val.shape}")    
    except:
        print(f"{key}: {val}")

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print(f"Train Size: {len(train)}")
print(f"Validation Size: {len(validate)}")
if test:
    print(f"Test Size: {len(test)}")

# CONFIG
LATENT_DIM = 64
HIDDEN_DIMS = [32, 64, 128, 256, 512]

NUM_EPOCHS = 50
BATCH_SIZE = 16
LEARNING_RATE = 3e-4

train_loader = DataLoader(train, batch_size=BATCH_SIZE, shuffle=True)
validate_loader = DataLoader(validate, batch_size=BATCH_SIZE, shuffle=False)
if test:
    test_loader = DataLoader(test, batch_size=BATCH_SIZE, shuffle=False)

model = VAE(
    audio_processor=train.audio_processor,
    latent_dim=LATENT_DIM,
    hidden_dims=HIDDEN_DIMS
).to(DEVICE)

def get_trainable_params(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Trainable parameters: {get_trainable_params(model)}")

optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

# Setup LR Scheduler
def lr_lambda(epoch):
    if epoch < 10:
        return 0.0003
    if epoch < 20:
        return 0.0002
    else:
        return 0.0001

scheduler = LambdaLR(optimizer, lr_lambda)

def loss_fn(x, x_hat, mu, log_var):
    # Measuring reconstruction of spectrogram and audio features
    spectrogram_loss =  F.mse_loss(x_hat['spectrogram'], x['spectrogram'], reduction="sum")

    family_loss =       F.mse_loss(x_hat['family'],      x['family'],      reduction="mean")
    instrument_loss =   F.mse_loss(x_hat['instrument'],  x['instrument'],  reduction="mean")
    source_loss =       F.mse_loss(x_hat['source'],      x['source'],      reduction="mean")
    note_loss =         F.mse_loss(x_hat['note'],        x['note'],        reduction="mean")
    qualities_loss =    F.mse_loss(x_hat['qualities'],   x['qualities'],   reduction="mean")
    id_loss =           F.mse_loss(x_hat['id'],          x['id'],          reduction="mean")

    # Measuring KL divergence
    kld_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
    
    total_loss = spectrogram_loss + family_loss + instrument_loss + source_loss + note_loss + qualities_loss + id_loss + kld_loss
 
    return total_loss

best_loss = float('inf')
best_model = None

# Training
for epoch in range(NUM_EPOCHS):
    
    # Train
    model.train()
    loop = tqdm(enumerate(train_loader))
    for i, x in loop:
        # Move tensors to device
        for key in x:
            x[key] = x[key].to(DEVICE)

        x_hat, mu, logvar = model(x)
        loss = loss_fn(x, x_hat, mu, logvar)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step() 
        loop.set_postfix(loss=loss.item())

    # Validate
    model.eval()
    with torch.no_grad():
        loss = 0

        for x in validate_loader:
            # Move tensors to device
            for key in x:
                x[key] = x[key].to(DEVICE)

            x_hat, mu, logvar = model(x)
            loss += loss_fn(x, x_hat, mu, logvar, epoch, 'validation')

        loss /= len(validate_loader)
        print(f"Epoch {epoch+1}/{NUM_EPOCHS} - Validation Loss: {loss}")

    # torch.save(model.state_dict(), f"./weights/samplegen_{LATENT_DIM}_{epoch+1}of{NUM_EPOCHS}_{BATCH_SIZE}_{LEARNING_RATE}.pt")
    if loss < best_loss:
        print("New best model found!")
        best_loss = loss
        best_model = model
    
    scheduler.step()

torch.save(best_model.state_dict(), f"./weights/best_samplegen_{LATENT_DIM}_{NUM_EPOCHS}_{BATCH_SIZE}_{LEARNING_RATE}.pt")
print("Training complete!")

if test:
    # Test model
    model.eval()
    with torch.no_grad():
        loss = 0

        for x in test_loader:
            # Move tensors to device
            for key in x:
                x[key] = x[key].to(DEVICE)

            x_hat, mu, logvar = best_model(x)
            loss += loss_fn(x, x_hat, mu, logvar, epoch, 'test')

        loss /= len(test_loader)
        print(f"Test Loss: {loss}")