train.py

import os
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
import gc  # For memory cleanup
import signal  # For graceful shutdown
import time  # For timestamp
from datetime import datetime
from models.model import create_model
from utils.data_loader import get_dataloaders
from utils.visualization import plot_training_history

# Global variables for graceful shutdown
INTERRUPT_FLAG = False
USERNAME = "madboy482"  # Current user's login

def handle_interrupt(signum, frame):
    """Handle interrupt signal (Ctrl+C) gracefully"""
    global INTERRUPT_FLAG
    print("\n\033[1;33mReceived interrupt signal. Will save model and exit after current epoch...\033[0m")
    INTERRUPT_FLAG = True

# Set up signal handler
signal.signal(signal.SIGINT, handle_interrupt)

# Custom tqdm class with colored bars
class ColorTqdm(tqdm):
    def __init__(self, *args, **kwargs):
        # Define color formatting
        bar_format = "{desc}: {percentage:3.0f}%|{bar:30}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]"
        
        # Choose a different color based on the type of operation
        if 'desc' in kwargs and '[Train]' in kwargs['desc']:
            bar_format = "\033[1;36m" + bar_format + "\033[0m"  # Cyan for training
        elif 'desc' in kwargs and '[Val]' in kwargs['desc']:
            bar_format = "\033[1;35m" + bar_format + "\033[0m"  # Magenta for validation
        else:
            bar_format = "\033[1;33m" + bar_format + "\033[0m"  # Yellow for other
            
        kwargs['bar_format'] = bar_format
        super().__init__(*args, **kwargs)

def train_model(train_dir, test_dir, output_dir, batch_size=16, num_epochs=20, 
               learning_rate=0.001, model_name='efficientnet-b0'):
    """
    Train the emotion recognition model
    
    Args:
        train_dir: Directory containing processed training data
        test_dir: Directory containing processed test data
        output_dir: Directory to save model checkpoints
        batch_size: Batch size for training
        num_epochs: Number of training epochs
        learning_rate: Learning rate for the optimizer
        model_name: EfficientNet variant to use
    """
    print("\033[1;32m" + "="*80 + "\033[0m")
    print("\033[1;32m Setting up for training...\033[0m")
    print(f"\033[1;32m Started by: {USERNAME} at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\033[0m")
    print("\033[1;32m" + "="*80 + "\033[0m")
    
    start_time = time.time()
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # Set device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"\033[1;33mUsing device: {device}\033[0m")
    
    # Get data loaders
    train_loader, val_loader, test_loader = get_dataloaders(
        train_dir, test_dir, batch_size=batch_size
    )
    
    # Create model
    print(f"\033[1;33mCreating model: {model_name}\033[0m")
    model = create_model(num_classes=7, model_name=model_name)
    model = model.to(device)
    
    # Define loss function and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    # Learning rate scheduler
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=0.5, patience=3, verbose=True
    )
    
    # Training and validation history
    train_losses = []
    val_losses = []
    train_accs = []
    val_accs = []
    
    # Best validation accuracy for model saving
    best_val_acc = 0.0
    best_epoch = -1
    patience = 5
    patience_counter = 0
    
    # Save the best model's state_dict
    best_model_state = None
    
    # Training loop
    print(f"\033[1;32mStarting training for {num_epochs} epochs...\033[0m")
    for epoch in range(num_epochs):
        epoch_start_time = time.time()
        
        # Memory cleanup before each epoch
        if device.type == 'cpu':
            gc.collect()
        
        # Training phase
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0
        
        # Use colored tqdm for training
        train_loop = ColorTqdm(train_loader, 
                          desc=f"Epoch {epoch+1}/{num_epochs} [Train]",
                          ascii=' ▁▂▃▄▅▆▇█')  # Custom ASCII characters for the bar
        for inputs, labels in train_loop:
            inputs, labels = inputs.to(device), labels.to(device)
            
            # Zero the parameter gradients
            optimizer.zero_grad()
            
            # Forward pass
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            # Backward pass and optimize
            loss.backward()
            optimizer.step()
            
            # Statistics
            train_loss += loss.item() * inputs.size(0)
            _, predicted = torch.max(outputs, 1)
            train_total += labels.size(0)
            train_correct += (predicted == labels).sum().item()
            
            # Update progress bar
            train_loop.set_postfix(loss=f"\033[1;36m{loss.item():.4f}\033[0m", 
                                  accuracy=f"\033[1;32m{train_correct/train_total:.4f}\033[0m")
        
        # Calculate training metrics
        train_loss = train_loss / len(train_loader.dataset)
        train_acc = train_correct / train_total
        train_losses.append(train_loss)
        train_accs.append(train_acc)
        
        # Validation phase
        model.eval()
        val_loss = 0.0
        val_correct = 0
        val_total = 0
        
        with torch.no_grad():
            # Use colored tqdm for validation
            val_loop = ColorTqdm(val_loader, 
                            desc=f"Epoch {epoch+1}/{num_epochs} [Val]",
                            ascii=' ▁▂▃▄▅▆▇█')
            for inputs, labels in val_loop:
                inputs, labels = inputs.to(device), labels.to(device)
                
                # Forward pass
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                
                # Statistics
                val_loss += loss.item() * inputs.size(0)
                _, predicted = torch.max(outputs, 1)
                val_total += labels.size(0)
                val_correct += (predicted == labels).sum().item()
                
                # Update progress bar
                val_loop.set_postfix(loss=f"\033[1;35m{loss.item():.4f}\033[0m", 
                                    accuracy=f"\033[1;32m{val_correct/val_total:.4f}\033[0m")
        
        # Calculate validation metrics
        val_loss = val_loss / len(val_loader.dataset)
        val_acc = val_correct / val_total
        val_losses.append(val_loss)
        val_accs.append(val_acc)
        
        # Update learning rate scheduler
        scheduler.step(val_loss)
        
        # Print epoch summary
        epoch_time = time.time() - epoch_start_time
        minutes, seconds = divmod(epoch_time, 60)
        
        print("\033[1;34m" + "-"*80 + "\033[0m")
        print(f"\033[1;34mEpoch {epoch+1}/{num_epochs} Summary:\033[0m")
        print(f"\033[1;32m  Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}\033[0m")
        print(f"\033[1;35m  Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}\033[0m")
        print(f"\033[1;33m  Epoch time: {int(minutes)}m {int(seconds)}s\033[0m")
        print(f"\033[1;33m  Current time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\033[0m")
        
        # Save best model (changed to save state_dict in memory)
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_epoch = epoch + 1
            # Save best model to disk
            best_model_path = os.path.join(output_dir, 'new_best_model_full.pth')
            torch.save(model.state_dict(), best_model_path)
            # Also keep a copy in memory
            best_model_state = model.state_dict().copy()
            print(f"\033[1;32m  ✓ Saved new best model with validation accuracy: {val_acc:.4f} (Epoch {best_epoch})\033[0m")
            patience_counter = 0
        else:
            patience_counter += 1
            print(f"\033[1;31m  ✗ No improvement. Patience: {patience_counter}/{patience}\033[0m")
            print(f"\033[1;33m    Best validation accuracy so far: {best_val_acc:.4f} (Epoch {best_epoch})\033[0m")
            
        # Save periodic checkpoint (every 3 epochs)
        if (epoch + 1) % 3 == 0:
            checkpoint_path = os.path.join(output_dir, f'checkpoint_epoch_{epoch+1}.pth')
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict(),
                'train_losses': train_losses,
                'val_losses': val_losses,
                'train_accs': train_accs,
                'val_accs': val_accs,
                'best_val_acc': best_val_acc,
                'best_epoch': best_epoch
            }, checkpoint_path)
            print(f"\033[1;36m  ⚙ Saved checkpoint at epoch {epoch+1}\033[0m")
        
        # Check for early stopping or interruption
        if patience_counter >= patience:
            print(f"\033[1;31mEarly stopping triggered after {epoch+1} epochs (no improvement for {patience} epochs)\033[0m")
            break
        
        if INTERRUPT_FLAG:
            print("\033[1;33mTraining interrupted by user. Saving current state...\033[0m")
            break
    
    # Calculate total training time
    total_time = time.time() - start_time
    hours, remainder = divmod(total_time, 3600)
    minutes, seconds = divmod(remainder, 60)
    
    print("\033[1;32m" + "="*80 + "\033[0m")
    print(f"\033[1;32mTraining completed!\033[0m")
    print(f"\033[1;33mTotal training time: {int(hours)} hours, {int(minutes)} minutes, {int(seconds)} seconds\033[0m")
    print(f"\033[1;33mFinished at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\033[0m")
    
    # Save final model (now using the best model state)
    final_model_path = os.path.join(output_dir, 'new_final_model_full.pth')
    
    # If we found a better model during training, use it as the final model
    if best_model_state is not None:
        torch.save(best_model_state, final_model_path)
        print(f"\033[1;32mFinal model saved - using best model from epoch {best_epoch} with validation accuracy: {best_val_acc:.4f}\033[0m")
    else:
        # This shouldn't normally happen, but as a fallback save the current model
        torch.save(model.state_dict(), final_model_path)
        print(f"\033[1;33mFinal model saved - using current model with validation accuracy: {val_acc:.4f}\033[0m")
    
    # Plot training history
    plot_training_history(train_losses, val_losses, train_accs, val_accs, 
                         best_epoch=best_epoch-1)  # Adjust for 0-indexing
    print("\033[1;32mTraining history plot saved\033[0m")
    
    # Save training history
    history = {
        'train_loss': train_losses,
        'val_loss': val_losses,
        'train_acc': train_accs,
        'val_acc': val_accs,
        'best_epoch': best_epoch
    }
    np.save(os.path.join(output_dir, 'new_training_history.npy'), history)
    print("\033[1;32mTraining history data saved\033[0m")
    print("\033[1;32m" + "="*80 + "\033[0m")
    
    return model, history

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train emotion recognition model")
    parser.add_argument("--train_dir", type=str, default="data/processed/train",
                        help="Directory containing processed training data")
    parser.add_argument("--test_dir", type=str, default="data/processed/test",
                        help="Directory containing processed test data")
    parser.add_argument("--output_dir", type=str, default="models/trained_models",
                        help="Directory to save model checkpoints")
    parser.add_argument("--batch_size", type=int, default=8,
                        help="Batch size for training")
    parser.add_argument("--num_epochs", type=int, default=15,
                        help="Number of training epochs")
    parser.add_argument("--learning_rate", type=float, default=0.001,
                        help="Learning rate for the optimizer")
    parser.add_argument("--model_name", type=str, default="efficientnet-b0",
                        help="EfficientNet variant to use")
    
    args = parser.parse_args()
    
    print(f"\033[1;36mStarting training at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\033[0m")
    print(f"\033[1;36mConfiguration: batch_size={args.batch_size}, epochs={args.num_epochs}, model={args.model_name}\033[0m")
    
    try:
        # Train the model
        train_model(
            train_dir=args.train_dir,
            test_dir=args.test_dir,
            output_dir=args.output_dir,
            batch_size=args.batch_size,
            num_epochs=args.num_epochs,
            learning_rate=args.learning_rate,
            model_name=args.model_name
        )
    except Exception as e:
        print(f"\033[1;31mError during training: {e}\033[0m")
        import traceback
        traceback.print_exc()
    finally:
        print(f"\033[1;36mTraining session ended at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\033[0m")