""" Enterprise-grade training pipeline with full monitoring and validation """ import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader import yaml import json from pathlib import Path import time from datetime import datetime import sys # Add project root to path sys.path.insert(0, str(Path(__file__).parent.parent)) from models.base.mnist_cnn import MNIST_CNN from utils.dataset_utils import load_mnist, create_dataloaders from utils.model_utils import save_model, evaluate_model, update_registry from utils.logging_utils import setup_logger class ModelTrainer: """Complete training pipeline with monitoring""" def __init__(self, config_path="config/training_config.yaml"): # Load configuration with open(config_path, 'r') as f: self.config = yaml.safe_load(f) # Setup self.device = torch.device(self.config['device']) self.logger = setup_logger('trainer', 'reports/logs/training.log') # Set random seed for reproducibility torch.manual_seed(self.config.get('seed', 42)) if torch.cuda.is_available(): torch.cuda.manual_seed_all(self.config.get('seed', 42)) self.logger.info(f"Training configuration: {json.dumps(self.config, indent=2)}") def setup_data(self): """Setup data loaders""" self.logger.info("Setting up data loaders...") # Load dataset train_set, test_set = load_mnist( augment=self.config.get('augment', False) ) # Create dataloaders with validation split self.train_loader, self.val_loader, self.test_loader = create_dataloaders( train_set, test_set, batch_size=self.config['batch_size'], val_split=self.config.get('validation_split', 0.1) ) self.logger.info(f"Training samples: {len(self.train_loader.dataset)}") self.logger.info(f"Validation samples: {len(self.val_loader.dataset)}") self.logger.info(f"Test samples: {len(self.test_loader.dataset)}") def setup_model(self): """Initialize model, optimizer, scheduler""" self.logger.info("Initializing model...") # Model self.model = MNIST_CNN().to(self.device) # Loss function self.criterion = nn.CrossEntropyLoss() # Optimizer if self.config['optimizer'].lower() == 'adam': self.optimizer = optim.Adam( self.model.parameters(), lr=self.config['learning_rate'], weight_decay=self.config.get('weight_decay', 0) ) elif self.config['optimizer'].lower() == 'sgd': self.optimizer = optim.SGD( self.model.parameters(), lr=self.config['learning_rate'], momentum=0.9, weight_decay=self.config.get('weight_decay', 0) ) else: raise ValueError(f"Unknown optimizer: {self.config['optimizer']}") # Learning rate scheduler if self.config.get('scheduler', 'none').lower() == 'step': self.scheduler = optim.lr_scheduler.StepLR( self.optimizer, step_size=10, gamma=0.1 ) else: self.scheduler = None # Log model summary total_params = sum(p.numel() for p in self.model.parameters()) self.logger.info(f"Model initialized with {total_params:,} parameters") def train_epoch(self, epoch: int): """Train for one epoch""" self.model.train() total_loss = 0 correct = 0 total = 0 for batch_idx, (data, target) in enumerate(self.train_loader): data, target = data.to(self.device), target.to(self.device) # Forward pass self.optimizer.zero_grad() output = self.model(data) loss = self.criterion(output, target) # Backward pass loss.backward() self.optimizer.step() # Statistics total_loss += loss.item() pred = output.argmax(dim=1) correct += pred.eq(target).sum().item() total += target.size(0) # Logging if batch_idx % self.config.get('log_frequency', 10) == 0: self.logger.debug( f"Epoch {epoch} [{batch_idx}/{len(self.train_loader)}] " f"Loss: {loss.item():.4f}" ) avg_loss = total_loss / len(self.train_loader) accuracy = 100. * correct / total return avg_loss, accuracy def validate(self): """Validate on validation set""" metrics = evaluate_model(self.model, self.val_loader, self.device) return metrics def save_checkpoint(self, epoch: int, is_best: bool = False): """Save model checkpoint""" checkpoint = { 'epoch': epoch, 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'config': self.config } # Regular checkpoint checkpoint_path = Path(self.config['save_path']).parent / f"checkpoint_epoch_{epoch}.pth" torch.save(checkpoint, checkpoint_path) # Best model if is_best: best_path = Path(self.config['save_path']).parent / "best_model.pth" torch.save(checkpoint, best_path) self.logger.info(f"Best model saved to {best_path}") def train(self): """Main training loop""" self.logger.info("Starting training...") start_time = time.time() best_val_acc = 0.0 patience_counter = 0 for epoch in range(1, self.config['epochs'] + 1): epoch_start = time.time() # Training train_loss, train_acc = self.train_epoch(epoch) # Validation val_metrics = self.validate() # Update scheduler if self.scheduler is not None: self.scheduler.step() # Calculate epoch time epoch_time = time.time() - epoch_start # Log results self.logger.info( f"Epoch {epoch:03d}/{self.config['epochs']:03d} | " f"Time: {epoch_time:.1f}s | " f"Train Loss: {train_loss:.4f} | " f"Train Acc: {train_acc:.2f}% | " f"Val Acc: {val_metrics['accuracy']:.2f}%" ) # Save checkpoint if epoch % self.config.get('checkpoint_frequency', 1) == 0: self.save_checkpoint(epoch) # Early stopping if val_metrics['accuracy'] > best_val_acc: best_val_acc = val_metrics['accuracy'] self.save_checkpoint(epoch, is_best=True) patience_counter = 0 else: patience_counter += 1 if patience_counter >= self.config.get('early_stopping_patience', float('inf')): self.logger.info(f"Early stopping triggered at epoch {epoch}") break # Final evaluation self.logger.info("Training completed. Running final evaluation...") test_metrics = evaluate_model(self.model, self.test_loader, self.device) total_time = time.time() - start_time self.logger.info(f"Total training time: {total_time:.1f}s") self.logger.info(f"Test Accuracy: {test_metrics['accuracy']:.2f}%") # Save final model metadata = { 'training_config': self.config, 'test_accuracy': test_metrics['accuracy'], 'test_loss': test_metrics['loss'], 'training_time': total_time, 'final_epoch': epoch } save_model(self.model, self.config['save_path'], metadata) update_registry('mnist_cnn', self.config['save_path'], metadata) return test_metrics def main(): """Main entry point""" trainer = ModelTrainer() trainer.setup_data() trainer.setup_model() results = trainer.train() print("\n" + "="*50) print("TRAINING COMPLETED SUCCESSFULLY") print("="*50) print(f"Final Test Accuracy: {results['accuracy']:.2f}%") print(f"Model saved to: {trainer.config['save_path']}") print("="*50) if __name__ == "__main__": main()