train.py

from model import *
from dataloader import SoundCamDataset, read_all_filenames
import torch
from torch.utils.data import DataLoader
import lightning as L
from lightning.pytorch.loggers import WandbLogger
from lightning.pytorch.callbacks import ModelCheckpoint
import argparse
import random
import numpy as np
from collections import defaultdict
import os

# Parse command line arguments
parser = argparse.ArgumentParser(description='SoundCam Training with Lightning')
parser.add_argument('--lr', type=float, default=3e-5, help='Learning rate')
parser.add_argument('--epochs', type=int, default=50, help='Number of epochs')
parser.add_argument('--batch_size', type=int, default=8, help='Batch size for training')
parser.add_argument('--data_path', type=str, default='data_preprocessed/', help='Path to the data directory')
parser.add_argument('--model', type=str, default='soundcam', help='Model to use')
args = parser.parse_args()

# Initialize wandb logger
wandb_logger = WandbLogger(project=args.model, name=args.model)

# Define callbacks
checkpoint_callback = ModelCheckpoint(
    monitor='val_loss',
    dirpath='checkpoints/',
    filename=f'{args.model}'+'-{epoch:02d}-{val_loss:.4f}',
    save_top_k=1,
    mode='min',
)

def set_seed(seed: int):
    random.seed(seed)  # Python random module
    np.random.seed(seed)  # NumPy
    torch.manual_seed(seed)  # PyTorch (CPU)
    torch.cuda.manual_seed(seed)  # PyTorch (single GPU)
    torch.cuda.manual_seed_all(seed)  # PyTorch (all GPUs)
    torch.backends.cudnn.deterministic = True  # Ensures reproducibility
    torch.backends.cudnn.benchmark = False  # Disables auto-optimizations

set_seed(42)

def train_test_split(filenames):
    train_filenames = []
    test_filenames = []
    
    # Group filenames by their directory
    files_by_dir = defaultdict(list)
    for filename in filenames:
        dir_name = os.path.dirname(filename)
        files_by_dir[dir_name].append(filename)
    
    # Split each directory's files into train and test
    for dir_name, files in files_by_dir.items():
        files.sort()  # Ensure files are sorted
        if files:
            test_filenames.append(files[-1])  # Last file for testing
            train_filenames.extend(files[:-1])  # All other files for training
    
    return train_filenames, test_filenames


# Prepare data
path = args.data_path
filenames = read_all_filenames(path)
train_filenames, test_filenames = train_test_split(filenames)

data = np.load(train_filenames[0])
spectrogram_time = data['fft'].shape[1]
spectrogram_channels = data['fft'].shape[2]-70

# Initialize model
if 'soundcam_swipe' in args.model:
    model = SoundCamModel_Swipe(learning_rate=args.lr, spectrogram_time=spectrogram_time, spectrogram_channels=spectrogram_channels)
elif 'baseline_swipe' in args.model:
    model = BaselineModel_Swipe(learning_rate=args.lr, spectrogram_time=spectrogram_time, spectrogram_channels=spectrogram_channels)
elif 'soundcam_widget' in args.model:
    model = SoundCamModel_Widget(learning_rate=args.lr, spectrogram_time=spectrogram_time, spectrogram_channels=spectrogram_channels)
elif 'baseline_widget' in args.model:
    model = BaselineModel_Widget(learning_rate=args.lr, spectrogram_time=spectrogram_time, spectrogram_channels=spectrogram_channels)
elif 'soundcam_object' in args.model:
    model = SoundCamModel_Object(learning_rate=args.lr, spectrogram_time=spectrogram_time, spectrogram_channels=spectrogram_channels)
elif 'baseline_object' in args.model:
    model = BaselineModel_Object(learning_rate=args.lr, spectrogram_time=spectrogram_time, spectrogram_channels=spectrogram_channels)

train_dataset = SoundCamDataset(train_filenames, model=args.model)
test_dataset = SoundCamDataset(test_filenames, model=args.model)

# open an ipython tin

train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=7)
val_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=7)

# Initialize trainer with gradient clipping
trainer = L.Trainer(
    max_epochs=args.epochs,
    logger=wandb_logger,
    callbacks=[checkpoint_callback],
    accelerator="gpu",
    devices=[0],
    log_every_n_steps=10,
    fast_dev_run=False,
    gradient_clip_val=1.0,  # Add gradient clipping
    gradient_clip_algorithm="norm"
)

# Log hyperparameters
wandb_logger.log_hyperparams({
    "learning_rate": args.lr,
    "epochs": args.epochs,
    "batch_size": args.batch_size,
})

# Train the model
trainer.fit(model, train_loader, val_loader)

# Get the path to the best model checkpoint
best_model_path = checkpoint_callback.best_model_path
print(f"Best model saved at: {best_model_path}")

# Load the best model
if 'soundcam_swipe' in args.model:
    best_model = SoundCamModel_Swipe.load_from_checkpoint(best_model_path, spectrogram_channels=spectrogram_channels, spectrogram_time=spectrogram_time)
elif 'baseline_swipe' in args.model:
    best_model = BaselineModel_Swipe.load_from_checkpoint(best_model_path, spectrogram_channels=spectrogram_channels, spectrogram_time=spectrogram_time)
elif 'soundcam_tap' in args.model:
    best_model = SoundCamModel_Tap.load_from_checkpoint(best_model_path, spectrogram_channels=spectrogram_channels, spectrogram_time=spectrogram_time)
elif 'baseline_tap' in args.model:
    best_model = BaselineModel_Tap.load_from_checkpoint(best_model_path, spectrogram_channels=spectrogram_channels, spectrogram_time=spectrogram_time)
elif 'soundcam_widget' in args.model:
    best_model = SoundCamModel_Widget.load_from_checkpoint(best_model_path, spectrogram_channels=spectrogram_channels, spectrogram_time=spectrogram_time)
elif 'baseline_widget' in args.model:
    best_model = BaselineModel_Widget.load_from_checkpoint(best_model_path, spectrogram_channels=spectrogram_channels, spectrogram_time=spectrogram_time)
elif 'soundcam_micro' in args.model:
    best_model = SoundCamModel_Micro.load_from_checkpoint(best_model_path, spectrogram_channels=spectrogram_channels, spectrogram_time=spectrogram_time)
elif 'baseline_micro' in args.model:
    best_model = BaselineModel_Micro.load_from_checkpoint(best_model_path, spectrogram_channels=spectrogram_channels, spectrogram_time=spectrogram_time)
elif 'soundcam_object' in args.model:
    best_model = SoundCamModel_Object.load_from_checkpoint(best_model_path, spectrogram_channels=spectrogram_channels, spectrogram_time=spectrogram_time)
elif 'baseline_object' in args.model:
    best_model = BaselineModel_Object.load_from_checkpoint(best_model_path, spectrogram_channels=spectrogram_channels, spectrogram_time=spectrogram_time)

# Save in PyTorch format
torch.save(best_model.state_dict(), "best_model.pth")
print("Best model saved as best_model.pth")

# Test the model with the best checkpoint
trainer.test(ckpt_path="best", dataloaders=val_loader)


print("Training complete!")