Verify Training Loop

Inigo_Parra: <p>Hi!<img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=12" title=":hugs:" class="emoji" alt=":hugs:" loading="lazy" width="20" height="20"></p>
<p>I am using Mel Spectrograms to classify sounds (24 classes). My training loop looks like this but I would like someone to verify if I am doing it correctly or if there are any issues that may be penalizing the model’s performance.</p>
<p>Also, what accuracy metric would be the best to judge my model? Standard or other type?</p>
<p>Here’s the code:</p>
<pre><code class="lang-auto">import torch
import torchaudio
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
from torch.nn.utils import clip_grad_norm_

import numpy as np
import random
import yaml
import os

from vit import VisionTransformer
from tools.optim_selector import set_optimizer
from tools.scheduler_selector import set_scheduler
from data import AudioData

import wandb


# For reproducibility, set the seed for all random number generators
def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(seed)
    random.seed(seed)

set_seed(42)


def save_checkpoint(model, optimizer, scheduler, epoch, path):
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict()
    }, path)


# TRAINING
def train(
        n_epochs: int, 
        model: nn.Module, 
        train_dataloader: DataLoader, 
        val_dataloader: DataLoader, 
        criterion: nn.Module, 
        optimizer: optim.Optimizer, 
        scheduler: optim.lr_scheduler, 
        device: torch.device, 
        wandb: bool = False,
        checkpoint_dir: str = 'checkpoints',
        checkpoint_interval: int = 20
    ):
    
    print(f"{'-'*50}\nDevice: {device}")
    print(f"Scheduler: {type(scheduler).__name__}\n{'-'*50}")
    print(f"Training...")
    
    model.to(device)
    if wandb:
        global_step = 0
        log_interval = 10

    # Make a checkpoint directory
    os.makedirs(checkpoint_dir, exist_ok=True)

    for epoch in range(n_epochs):
        # TRAIN
        model.train()
        running_train_loss = 0.0
        correct_train = 0
        total_train = 0
        for batch_idx, (signals, labels) in enumerate(train_dataloader):
            signals, labels = signals.to(device), labels.to(device)
            
            # expected signals shape should be [batch_size, channels, height, width]
            if len(signals.shape) != 4:
                signals = signals.unsqueeze(1)
            
            outputs = model(signals)
            loss = criterion(outputs, labels)

            optimizer.zero_grad()
            loss.backward()
            clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

            running_train_loss += loss.item()

            _, predicted = torch.max(outputs.data, 1)
            total_train += labels.size(0)
            correct_train += (predicted == labels).sum().item()

            if wandb:
                global_step += 1

            # Print step metrics in the local console
            if batch_idx % 10 == 0:
                print(f'Epoch [{epoch+1}/{n_epochs}] - Step [{batch_idx+1}/{len(train_dataloader)}] - Loss: {loss.item():.3f}')

            train_accuracy = (correct_train / total_train) * 100

            # Log metrics to wandb
            if wandb and global_step % log_interval == 0:
                wandb.log({
                    'step': global_step,
                    'train_loss': loss.item(),
                    'train_accuracy': train_accuracy,
                    'learning_rate': scheduler.get_last_lr()
                })

        epoch_train_loss = running_train_loss / len(train_dataloader)
        # Print epoch metrics in the local console
        print(f'Epoch [{epoch+1}/{n_epochs}] - Train Loss: {epoch_train_loss:.3f} || Acc: {train_accuracy:.3f}')


        # VALIDATION
        model.eval()
        running_val_loss = 0.0
        correct = 0
        total = 0
        with torch.no_grad():
            for signals, labels in val_dataloader:
                signals, labels = signals.to(device), labels.to(device)
                
                if len(signals.shape) == 4:
                    signals = signals.squeeze(1)
                
                signals = signals.unsqueeze(1)
                
                outputs = model(signals)
                loss = criterion(outputs, labels)
                running_val_loss += loss.item()

                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        epoch_val_loss = running_val_loss / len(val_dataloader)
        val_accuracy = (correct / total) * 100

        # Pass loss to scheduler and update learning rate (if needed)
        if scheduler is not None:
            scheduler.step()

        #Log validation metrics to wandb
        if wandb:
            wandb.log({
                'step': global_step,
                'val_loss': epoch_val_loss,
                'val_accuracy': val_accuracy
            })

        # Print LR and summary
        print(f'Learning rate: {scheduler.get_last_lr()}')
        print(f'Epoch [{epoch+1}/{n_epochs}] - Train Loss: {epoch_train_loss:.3f} - Val Loss: {epoch_val_loss:.3f} || Val Accuracy: {val_accuracy:.3f}')

        # Save checkpoint every x epochs
        if epoch % checkpoint_interval == 0 and epoch != 0:
            checkpoint_path = os.path.join(checkpoint_dir, f'checkpoint_{epoch+1}.pt')
            save_checkpoint(model, optimizer, scheduler, epoch, checkpoint_path)

    print("Training complete.")


# EVALUATION IN TEST SET
def evaluate(model: nn.Module, test_dataloader: DataLoader, criterion: nn.Module, device: torch.device):
    print("Evaluating...")
    model.to(device)
    model.eval()
    test_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        for signals, labels in test_dataloader:
            signals, labels = signals.to(device), labels.to(device)

            if len(signals.shape) == 4:
                signals = signals.squeeze(1)

            signals = signals.unsqueeze(1)

            outputs = model(signals)
            loss = criterion(outputs, labels)
            test_loss += loss.item()

            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    test_loss = test_loss / len(test_dataloader)
    test_accuracy = (correct / total) * 100
    
    # Evaluation results
    print(f'Test Loss: {test_loss:.3f} || Test Accuracy: {test_accuracy:.3f}')
    print("Evaluation complete.")
</code></pre>
<p>I can share the repo if needed too.</p>
<p>Thank you so much!<img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=12" title=":hugs:" class="emoji" alt=":hugs:" loading="lazy" width="20" height="20"></p>