""" Base trainer classes and interfaces for the Trixy ML training framework. This module provides the fundamental building blocks for all ML trainers, including abstract base classes, configuration management, and training state tracking. """ import os import time import logging import hashlib from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional, Union, Callable, Tuple from dataclasses import dataclass, field from enum import Enum import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader import numpy as np class TrainingState(Enum): """Training state enumeration for tracking trainer status.""" IDLE = "idle" PREPARING = "preparing" TRAINING = "training" VALIDATING = "validating" TESTING = "testing" PAUSED = "paused" COMPLETED = "completed" FAILED = "failed" STOPPING = "stopping" class ModelFormat(Enum): """Supported model formats.""" PTH = ".pth" PT = ".pt" ONNX = ".onnx" @dataclass class TrainerConfig: """Configuration class for ML trainers.""" # General settings trainer_name: str = "BaseTrainer" model_name: str = "model" output_dir: str = "./models" data_dir: str = "./trainer/data" config_dir: str = "./config" # Training parameters batch_size: int = 32 learning_rate: float = 0.001 num_epochs: int = 100 min_epochs: int = 10 max_epochs: int = 1000 early_stopping_patience: int = 10 # Model parameters model_format: ModelFormat = ModelFormat.PTH use_mixed_precision: bool = True gradient_clip_norm: float = 1.0 weight_decay: float = 1e-4 # Data parameters validation_split: float = 0.2 test_split: float = 0.1 shuffle_data: bool = True num_workers: int = 4 pin_memory: bool = True # Audio processing parameters sample_rate: int = 16000 bit_depth: int = 16 audio_length: float = 1.5 # seconds n_mels: int = 40 n_fft: int = 512 hop_length: int = 160 # 10ms at 16kHz win_length: int = 400 # 25ms at 16kHz # Augmentation parameters use_augmentation: bool = True noise_factor: float = 0.1 speed_factor: float = 0.1 pitch_factor: float = 0.1 volume_factor: float = 0.2 # Security settings use_password_protection: bool = True password: Optional[str] = None encryption_algorithm: str = "AES256" # Logging and monitoring log_level: str = "INFO" save_checkpoints: bool = True checkpoint_interval: int = 10 # epochs validate_interval: int = 1 # epochs log_interval: int = 100 # batches # Hardware settings device: str = "auto" # auto, cpu, cuda, mps use_distributed: bool = False num_gpus: int = 1 # Advanced settings resume_from_checkpoint: bool = False checkpoint_path: Optional[str] = None freeze_backbone: bool = False use_scheduler: bool = True scheduler_type: str = "cosine" # cosine, step, exponential # Custom parameters (extensible) custom_params: Dict[str, Any] = field(default_factory=dict) def __post_init__(self): """Validate and process configuration after initialization.""" self._validate_config() self._process_device() def _validate_config(self): """Validate configuration parameters.""" if self.batch_size <= 0: raise ValueError("batch_size must be positive") if self.learning_rate <= 0: raise ValueError("learning_rate must be positive") if self.num_epochs <= 0: raise ValueError("num_epochs must be positive") if not 0 < self.validation_split < 1: raise ValueError("validation_split must be between 0 and 1") if not 0 <= self.test_split < 1: raise ValueError("test_split must be between 0 and 1") if self.validation_split + self.test_split >= 1: raise ValueError("validation_split + test_split must be less than 1") def _process_device(self): """Automatically determine the best available device.""" if self.device == "auto": if torch.cuda.is_available(): self.device = "cuda" elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): self.device = "mps" else: self.device = "cpu" def to_dict(self) -> Dict[str, Any]: """Convert configuration to dictionary.""" config_dict = {} for key, value in self.__dict__.items(): if isinstance(value, Enum): config_dict[key] = value.value else: config_dict[key] = value return config_dict @classmethod def from_dict(cls, config_dict: Dict[str, Any]) -> 'TrainerConfig': """Create configuration from dictionary.""" # Handle enum conversion if 'model_format' in config_dict: config_dict['model_format'] = ModelFormat(config_dict['model_format']) return cls(**config_dict) class TrainingMetrics: """Class for tracking and managing training metrics.""" def __init__(self): self.metrics = { 'train_loss': [], 'train_accuracy': [], 'val_loss': [], 'val_accuracy': [], 'learning_rates': [], 'epoch_times': [], 'total_time': 0.0 } self.current_epoch = 0 self.best_val_loss = float('inf') self.best_val_accuracy = 0.0 self.best_epoch = 0 def update(self, epoch: int, train_loss: float, train_acc: float = None, val_loss: float = None, val_acc: float = None, lr: float = None, epoch_time: float = None): """Update metrics for the current epoch.""" self.current_epoch = epoch self.metrics['train_loss'].append(train_loss) if train_acc is not None: self.metrics['train_accuracy'].append(train_acc) if val_loss is not None: self.metrics['val_loss'].append(val_loss) if val_loss < self.best_val_loss: self.best_val_loss = val_loss self.best_epoch = epoch if val_acc is not None: self.metrics['val_accuracy'].append(val_acc) if val_acc > self.best_val_accuracy: self.best_val_accuracy = val_acc if lr is not None: self.metrics['learning_rates'].append(lr) if epoch_time is not None: self.metrics['epoch_times'].append(epoch_time) self.metrics['total_time'] += epoch_time def get_summary(self) -> Dict[str, Any]: """Get a summary of training metrics.""" return { 'current_epoch': self.current_epoch, 'best_epoch': self.best_epoch, 'best_val_loss': self.best_val_loss, 'best_val_accuracy': self.best_val_accuracy, 'total_training_time': self.metrics['total_time'], 'avg_epoch_time': np.mean(self.metrics['epoch_times']) if self.metrics['epoch_times'] else 0, 'final_train_loss': self.metrics['train_loss'][-1] if self.metrics['train_loss'] else None, 'final_val_loss': self.metrics['val_loss'][-1] if self.metrics['val_loss'] else None } class BaseTrainer(ABC): """ Abstract base class for all ML trainers in the Trixy framework. This class provides the fundamental structure and common functionality for all training implementations, including state management, logging, checkpointing, and validation. """ def __init__(self, config: TrainerConfig): """ Initialize the base trainer. Args: config: TrainerConfig object containing training parameters """ self.config = config self.state = TrainingState.IDLE self.metrics = TrainingMetrics() self.logger = self._setup_logger() # Training components (to be initialized by subclasses) self.model: Optional[nn.Module] = None self.optimizer: Optional[optim.Optimizer] = None self.scheduler: Optional[optim.lr_scheduler._LRScheduler] = None self.criterion: Optional[nn.Module] = None self.scaler: Optional[torch.cuda.amp.GradScaler] = None # Data loaders (to be initialized by subclasses) self.train_loader: Optional[DataLoader] = None self.val_loader: Optional[DataLoader] = None self.test_loader: Optional[DataLoader] = None # Training state self.current_epoch = 0 self.global_step = 0 self.best_model_state = None self.early_stopping_counter = 0 # Callbacks self.callbacks: List[Callable] = [] # Setup mixed precision if requested if self.config.use_mixed_precision and self.config.device != "cpu": self.scaler = torch.cuda.amp.GradScaler() self.logger.info(f"Initialized {self.__class__.__name__} with config: {self.config.trainer_name}") def _setup_logger(self) -> logging.Logger: """Setup logger for the trainer.""" logger = logging.getLogger(f"trixy.trainer.{self.config.trainer_name}") logger.setLevel(getattr(logging, self.config.log_level.upper())) if not logger.handlers: # Create console handler console_handler = logging.StreamHandler() console_handler.setLevel(logging.INFO) # Create file handler os.makedirs(f"{self.config.output_dir}/logs", exist_ok=True) file_handler = logging.FileHandler( f"{self.config.output_dir}/logs/{self.config.trainer_name}.log" ) file_handler.setLevel(logging.DEBUG) # Create formatter formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) console_handler.setFormatter(formatter) file_handler.setFormatter(formatter) logger.addHandler(console_handler) logger.addHandler(file_handler) return logger @abstractmethod def prepare_data(self) -> Tuple[DataLoader, DataLoader, DataLoader]: """ Prepare training, validation, and test data loaders. Returns: Tuple of (train_loader, val_loader, test_loader) """ pass @abstractmethod def build_model(self) -> nn.Module: """ Build and return the model architecture. Returns: PyTorch model """ pass @abstractmethod def create_criterion(self) -> nn.Module: """ Create and return the loss criterion. Returns: PyTorch loss function """ pass def create_optimizer(self) -> optim.Optimizer: """Create and return the optimizer.""" return optim.AdamW( self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay ) def create_scheduler(self) -> Optional[optim.lr_scheduler._LRScheduler]: """Create and return the learning rate scheduler.""" if not self.config.use_scheduler: return None if self.config.scheduler_type == "cosine": return optim.lr_scheduler.CosineAnnealingLR( self.optimizer, T_max=self.config.num_epochs ) elif self.config.scheduler_type == "step": return optim.lr_scheduler.StepLR( self.optimizer, step_size=30, gamma=0.1 ) elif self.config.scheduler_type == "exponential": return optim.lr_scheduler.ExponentialLR( self.optimizer, gamma=0.95 ) else: self.logger.warning(f"Unknown scheduler type: {self.config.scheduler_type}") return None def setup_training(self): """Setup all training components.""" self.state = TrainingState.PREPARING self.logger.info("Setting up training components...") # Prepare data self.train_loader, self.val_loader, self.test_loader = self.prepare_data() # Build model self.model = self.build_model() self.model.to(self.config.device) # Create training components self.criterion = self.create_criterion() self.optimizer = self.create_optimizer() self.scheduler = self.create_scheduler() # Load checkpoint if requested if self.config.resume_from_checkpoint and self.config.checkpoint_path: self.load_checkpoint(self.config.checkpoint_path) self.logger.info("Training setup completed successfully") def train_epoch(self) -> Tuple[float, float]: """ Train for one epoch. Returns: Tuple of (average_loss, average_accuracy) """ self.model.train() total_loss = 0.0 total_correct = 0 total_samples = 0 for batch_idx, (data, targets) in enumerate(self.train_loader): data = data.to(self.config.device) targets = targets.to(self.config.device) # Zero gradients self.optimizer.zero_grad() # Forward pass with mixed precision if self.scaler is not None: with torch.cuda.amp.autocast(): outputs = self.model(data) loss = self.criterion(outputs, targets) # Backward pass self.scaler.scale(loss).backward() # Gradient clipping if self.config.gradient_clip_norm > 0: self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_( self.model.parameters(), self.config.gradient_clip_norm ) self.scaler.step(self.optimizer) self.scaler.update() else: outputs = self.model(data) loss = self.criterion(outputs, targets) loss.backward() # Gradient clipping if self.config.gradient_clip_norm > 0: torch.nn.utils.clip_grad_norm_( self.model.parameters(), self.config.gradient_clip_norm ) self.optimizer.step() # Update metrics total_loss += loss.item() if len(outputs.shape) > 1 and outputs.shape[1] > 1: # Classification _, predicted = torch.max(outputs.data, 1) total_correct += (predicted == targets).sum().item() total_samples += targets.size(0) self.global_step += 1 # Log progress if batch_idx % self.config.log_interval == 0: self.logger.debug( f"Epoch {self.current_epoch}, Batch {batch_idx}/{len(self.train_loader)}, " f"Loss: {loss.item():.6f}" ) avg_loss = total_loss / len(self.train_loader) avg_accuracy = total_correct / total_samples if total_samples > 0 else 0.0 return avg_loss, avg_accuracy def validate_epoch(self) -> Tuple[float, float]: """ Validate for one epoch. Returns: Tuple of (average_loss, average_accuracy) """ self.model.eval() total_loss = 0.0 total_correct = 0 total_samples = 0 with torch.no_grad(): for data, targets in self.val_loader: data = data.to(self.config.device) targets = targets.to(self.config.device) outputs = self.model(data) loss = self.criterion(outputs, targets) total_loss += loss.item() if len(outputs.shape) > 1 and outputs.shape[1] > 1: # Classification _, predicted = torch.max(outputs.data, 1) total_correct += (predicted == targets).sum().item() total_samples += targets.size(0) avg_loss = total_loss / len(self.val_loader) avg_accuracy = total_correct / total_samples if total_samples > 0 else 0.0 return avg_loss, avg_accuracy def train(self) -> TrainingMetrics: """ Main training loop. Returns: TrainingMetrics object containing training history """ try: self.setup_training() self.state = TrainingState.TRAINING self.logger.info(f"Starting training for {self.config.num_epochs} epochs") start_time = time.time() for epoch in range(self.current_epoch, self.config.num_epochs): epoch_start_time = time.time() self.current_epoch = epoch # Training phase train_loss, train_acc = self.train_epoch() # Validation phase if epoch % self.config.validate_interval == 0: self.state = TrainingState.VALIDATING val_loss, val_acc = self.validate_epoch() self.state = TrainingState.TRAINING else: val_loss, val_acc = None, None # Update learning rate if self.scheduler is not None: self.scheduler.step() # Calculate epoch time epoch_time = time.time() - epoch_start_time # Update metrics current_lr = self.optimizer.param_groups[0]['lr'] self.metrics.update( epoch=epoch, train_loss=train_loss, train_acc=train_acc, val_loss=val_loss, val_acc=val_acc, lr=current_lr, epoch_time=epoch_time ) # Log progress log_msg = f"Epoch {epoch}/{self.config.num_epochs} - " log_msg += f"Train Loss: {train_loss:.6f}, Train Acc: {train_acc:.4f}" if val_loss is not None: log_msg += f", Val Loss: {val_loss:.6f}, Val Acc: {val_acc:.4f}" log_msg += f", LR: {current_lr:.2e}, Time: {epoch_time:.2f}s" self.logger.info(log_msg) # Save checkpoint if self.config.save_checkpoints and epoch % self.config.checkpoint_interval == 0: self.save_checkpoint(epoch) # Early stopping check if val_loss is not None: if val_loss < self.metrics.best_val_loss: self.best_model_state = self.model.state_dict().copy() self.early_stopping_counter = 0 else: self.early_stopping_counter += 1 if self.early_stopping_counter >= self.config.early_stopping_patience: self.logger.info(f"Early stopping triggered after {epoch + 1} epochs") break # Check minimum epochs if epoch < self.config.min_epochs: continue # Execute callbacks for callback in self.callbacks: callback(self, epoch) total_time = time.time() - start_time self.logger.info(f"Training completed in {total_time:.2f} seconds") # Load best model if available if self.best_model_state is not None: self.model.load_state_dict(self.best_model_state) self.state = TrainingState.COMPLETED return self.metrics except Exception as e: self.state = TrainingState.FAILED self.logger.error(f"Training failed: {str(e)}") raise def test(self) -> Dict[str, float]: """ Test the trained model. Returns: Dictionary containing test metrics """ self.state = TrainingState.TESTING self.logger.info("Starting model testing...") if self.test_loader is None: raise ValueError("Test loader not available") self.model.eval() total_loss = 0.0 total_correct = 0 total_samples = 0 with torch.no_grad(): for data, targets in self.test_loader: data = data.to(self.config.device) targets = targets.to(self.config.device) outputs = self.model(data) loss = self.criterion(outputs, targets) total_loss += loss.item() if len(outputs.shape) > 1 and outputs.shape[1] > 1: # Classification _, predicted = torch.max(outputs.data, 1) total_correct += (predicted == targets).sum().item() total_samples += targets.size(0) test_loss = total_loss / len(self.test_loader) test_accuracy = total_correct / total_samples if total_samples > 0 else 0.0 test_results = { 'test_loss': test_loss, 'test_accuracy': test_accuracy, 'total_samples': total_samples } self.logger.info(f"Test Results - Loss: {test_loss:.6f}, Accuracy: {test_accuracy:.4f}") return test_results def save_checkpoint(self, epoch: int, filepath: Optional[str] = None): """Save training checkpoint.""" if filepath is None: os.makedirs(f"{self.config.output_dir}/checkpoints", exist_ok=True) filepath = f"{self.config.output_dir}/checkpoints/{self.config.model_name}_epoch_{epoch}.pth" checkpoint = { 'epoch': epoch, 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None, 'scaler_state_dict': self.scaler.state_dict() if self.scaler else None, 'metrics': self.metrics, 'config': self.config.to_dict(), 'best_model_state': self.best_model_state } torch.save(checkpoint, filepath) self.logger.info(f"Checkpoint saved: {filepath}") def load_checkpoint(self, filepath: str): """Load training checkpoint.""" if not os.path.exists(filepath): raise FileNotFoundError(f"Checkpoint not found: {filepath}") checkpoint = torch.load(filepath, map_location=self.config.device) # Load model state if self.model is not None: self.model.load_state_dict(checkpoint['model_state_dict']) # Load optimizer state if self.optimizer is not None: self.optimizer.load_state_dict(checkpoint['optimizer_state_dict']) # Load scheduler state if self.scheduler is not None and checkpoint.get('scheduler_state_dict'): self.scheduler.load_state_dict(checkpoint['scheduler_state_dict']) # Load scaler state if self.scaler is not None and checkpoint.get('scaler_state_dict'): self.scaler.load_state_dict(checkpoint['scaler_state_dict']) # Load training state self.current_epoch = checkpoint['epoch'] + 1 self.metrics = checkpoint.get('metrics', TrainingMetrics()) self.best_model_state = checkpoint.get('best_model_state') self.logger.info(f"Checkpoint loaded: {filepath}") def add_callback(self, callback: Callable): """Add a callback function to be executed during training.""" self.callbacks.append(callback) def get_model_summary(self) -> str: """Get a summary of the model architecture.""" if self.model is None: return "Model not initialized" total_params = sum(p.numel() for p in self.model.parameters()) trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) summary = f"Model: {self.model.__class__.__name__}\n" summary += f"Total parameters: {total_params:,}\n" summary += f"Trainable parameters: {trainable_params:,}\n" summary += f"Device: {next(self.model.parameters()).device}\n" return summary def stop_training(self): """Stop training gracefully.""" self.state = TrainingState.STOPPING self.logger.info("Training stop requested") def pause_training(self): """Pause training.""" self.state = TrainingState.PAUSED self.logger.info("Training paused") def resume_training(self): """Resume training from paused state.""" if self.state == TrainingState.PAUSED: self.state = TrainingState.TRAINING self.logger.info("Training resumed") def get_training_status(self) -> Dict[str, Any]: """Get current training status and progress.""" return { 'state': self.state.value, 'current_epoch': self.current_epoch, 'total_epochs': self.config.num_epochs, 'progress_percentage': (self.current_epoch / self.config.num_epochs) * 100, 'global_step': self.global_step, 'metrics_summary': self.metrics.get_summary(), 'best_epoch': self.metrics.best_epoch, 'early_stopping_counter': self.early_stopping_counter }