Upload src/training\trainer.py with huggingface_hub

src/training//trainer.py

@@ -0,0 +1,414 @@
+"""
+Enhanced trainer for architectural style classification.
+Includes advanced optimization techniques for better accuracy.
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, OneCycleLR, ReduceLROnPlateau
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, LearningRateMonitor
+from pytorch_lightning.loggers import TensorBoardLogger
+import numpy as np
+from typing import Dict, List, Optional, Tuple, Any
+import os
+import json
+from datetime import datetime
+
+from .losses import HierarchicalLoss, ContrastiveLoss, StyleRelationshipLoss, FocalLoss, LabelSmoothingLoss
+from .metrics import ArchitecturalMetrics
+from .data_loader import EnhancedArchitecturalDataLoader
+
+
+class EnhancedArchitecturalTrainer(pl.LightningModule):
+    """Enhanced trainer for architectural style classification with advanced optimization."""
+    
+    def __init__(self, model: nn.Module, config: Dict[str, Any]):
+        super().__init__()
+        self.model = model
+        self.config = config
+        self.save_hyperparameters(ignore=['model'])
+        
+        # Enhanced configuration
+        self.learning_rate = config.get('learning_rate', 1e-4)
+        self.weight_decay = config.get('weight_decay', 1e-4)
+        self.batch_size = config.get('batch_size', 8)
+        self.num_classes = config.get('num_classes', 25)
+        self.use_mixed_precision = config.get('use_mixed_precision', True)
+        self.use_early_stopping = config.get('use_early_stopping', True)
+        self.patience = config.get('patience', 15)
+        self.gradient_clip_val = config.get('gradient_clip_val', 1.0)
+        self.accumulate_grad_batches = config.get('accumulate_grad_batches', 2)
+        
+        # Enhanced loss functions
+        self.use_focal_loss = config.get('use_focal_loss', True)
+        self.use_label_smoothing = config.get('use_label_smoothing', True)
+        self.use_contrastive_loss = config.get('use_contrastive_loss', True)
+        
+        # Initialize loss functions
+        self._init_loss_functions()
+        
+        # Initialize metrics
+        self.metrics = ArchitecturalMetrics(num_classes=self.num_classes)
+        
+        # Curriculum learning
+        self.curriculum_stage = 0
+        self.curriculum_classes_count = self.num_classes
+        
+        # Learning rate scheduling
+        self.scheduler_step_size = config.get('scheduler_step_size', 10)
+        self.scheduler_gamma = config.get('scheduler_gamma', 0.5)
+        self.warmup_epochs = config.get('warmup_epochs', 5)
+        
+        # TensorBoard logger
+        self.tensorboard_logger = TensorBoardLogger(
+            save_dir='logs',
+            name=f'architectural_training_{datetime.now().strftime("%Y%m%d_%H%M%S")}',
+            version=None
+        )
+    
+    def _init_loss_functions(self):
+        """Initialize enhanced loss functions."""
+        # Main classification loss
+        if self.use_focal_loss:
+            self.classification_loss = FocalLoss(
+                alpha=1.0, 
+                gamma=2.0, 
+                num_classes=self.num_classes
+            )
+        elif self.use_label_smoothing:
+            self.classification_loss = LabelSmoothingLoss(
+                smoothing=0.1, 
+                num_classes=self.num_classes
+            )
+        else:
+            self.classification_loss = nn.CrossEntropyLoss()
+        
+        # Additional loss functions
+        if self.use_contrastive_loss:
+            self.contrastive_loss = ContrastiveLoss(temperature=0.07)
+        
+        # Hierarchical loss for multi-scale features
+        self.hierarchical_loss = HierarchicalLoss(
+            num_classes=self.num_classes,
+            hierarchy_weights=[1.0, 0.5, 0.25]
+        )
+        
+        # Style relationship loss
+        self.style_relationship_loss = StyleRelationshipLoss(
+            num_classes=self.num_classes,
+            temperature=0.1
+        )
+    
+    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """Forward pass through the model."""
+        return self.model(x)
+    
+    def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> Dict[str, torch.Tensor]:
+        """Enhanced training step with multiple loss components."""
+        images, labels = batch
+        
+        # Forward pass
+        outputs = self(images)
+        
+        # Extract logits
+        if isinstance(outputs, dict):
+            logits = outputs.get('fine_logits', outputs.get('logits'))
+            features = outputs.get('features', None)
+            hierarchical_outputs = outputs.get('hierarchical_outputs', None)
+        else:
+            logits = outputs
+            features = None
+            hierarchical_outputs = None
+        
+        # Calculate main classification loss
+        if self.use_focal_loss or self.use_label_smoothing:
+            main_loss = self.classification_loss(logits, labels)
+        else:
+            main_loss = self.classification_loss(logits, labels)
+        
+        # Calculate additional losses
+        total_loss = main_loss
+        loss_dict = {'main_loss': main_loss}
+        
+        # Hierarchical loss
+        if hierarchical_outputs is not None:
+            hierarchical_loss = self.hierarchical_loss(hierarchical_outputs, labels)
+            total_loss += 0.3 * hierarchical_loss
+            loss_dict['hierarchical_loss'] = hierarchical_loss
+        
+        # Contrastive loss
+        if self.use_contrastive_loss and features is not None:
+            contrastive_loss = self.contrastive_loss(features, labels)
+            total_loss += 0.1 * contrastive_loss
+            loss_dict['contrastive_loss'] = contrastive_loss
+        
+        # Style relationship loss
+        style_loss = self.style_relationship_loss(logits, labels)
+        total_loss += 0.05 * style_loss
+        loss_dict['style_loss'] = style_loss
+        
+        # Calculate metrics
+        with torch.no_grad():
+            metrics = self.metrics.compute(logits, labels)
+            for key, value in metrics.items():
+                if isinstance(value, (int, float)):
+                    self.log(f'train_{key}', float(value), prog_bar=True)
+        
+        # Log losses
+        loss_dict['loss'] = total_loss
+        for key, value in loss_dict.items():
+            self.log(f'train_{key}', value, prog_bar=True)
+        
+        return loss_dict
+    
+    def validation_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> Dict[str, torch.Tensor]:
+        """Enhanced validation step."""
+        images, labels = batch
+        
+        # Forward pass
+        outputs = self(images)
+        
+        # Extract logits
+        if isinstance(outputs, dict):
+            logits = outputs.get('fine_logits', outputs.get('logits'))
+        else:
+            logits = outputs
+        
+        # Calculate loss
+        val_loss = self.classification_loss(logits, labels)
+        
+        # Calculate metrics
+        metrics = self.metrics.compute(logits, labels)
+        
+        # Log validation metrics
+        self.log('val_loss', val_loss, prog_bar=True)
+        for key, value in metrics.items():
+            if isinstance(value, (int, float)):
+                self.log(f'val_{key}', float(value), prog_bar=True)
+        
+        return {'val_loss': val_loss, 'logits': logits, 'labels': labels}
+    
+    def on_validation_epoch_end(self) -> None:
+        """Enhanced validation epoch end with detailed logging."""
+        # Log curriculum learning progress
+        self.log('curriculum_stage', float(self.curriculum_stage), prog_bar=True)
+        self.log('curriculum_classes_count', float(self.curriculum_classes_count), prog_bar=True)
+        
+        # Log learning rate
+        current_lr = self.optimizers().param_groups[0]['lr']
+        self.log('learning_rate', current_lr, prog_bar=True)
+    
+    def test_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> Dict[str, torch.Tensor]:
+        """Enhanced test step."""
+        images, labels = batch
+        
+        # Forward pass
+        outputs = self(images)
+        
+        # Extract logits
+        if isinstance(outputs, dict):
+            logits = outputs.get('fine_logits', outputs.get('logits'))
+        else:
+            logits = outputs
+        
+        # Calculate metrics
+        metrics = self.metrics.compute(logits, labels)
+        
+        # Log test metrics
+        for key, value in metrics.items():
+            if isinstance(value, (int, float)):
+                self.log(f'test_{key}', float(value), prog_bar=True)
+        
+        return {'logits': logits, 'labels': labels}
+    
+    def on_test_epoch_end(self) -> None:
+        """Save test results."""
+        # Save confusion matrix
+        confusion_matrix = self.metrics.confusion_matrix
+        if confusion_matrix is not None:
+            np.save('results/confusion_matrix.npy', confusion_matrix.cpu().numpy())
+        
+        # Save detailed results
+        results = {
+            'model_name': self.model.__class__.__name__,
+            'config': self.config,
+            'test_metrics': {
+                'accuracy': float(self.metrics.accuracy),
+                'precision_macro': float(self.metrics.precision_macro),
+                'recall_macro': float(self.metrics.recall_macro),
+                'f1_macro': float(self.metrics.f1_macro),
+                'precision_weighted': float(self.metrics.precision_weighted),
+                'recall_weighted': float(self.metrics.recall_weighted),
+                'f1_weighted': float(self.metrics.f1_weighted),
+            }
+        }
+        
+        # Save results
+        os.makedirs('results', exist_ok=True)
+        with open(f'results/{self.config.get("experiment_name", "test")}_results.json', 'w') as f:
+            json.dump(results, f, indent=2)
+    
+    def configure_optimizers(self):
+        """Configure enhanced optimizers and schedulers."""
+        # Enhanced optimizer with better parameters
+        optimizer = optim.AdamW(
+            self.parameters(),
+            lr=self.learning_rate,
+            weight_decay=self.weight_decay,
+            betas=(0.9, 0.999),
+            eps=1e-8
+        )
+        
+        # Enhanced learning rate scheduler
+        scheduler = CosineAnnealingWarmRestarts(
+            optimizer,
+            T_0=10,  # Restart every 10 epochs
+            T_mult=2,  # Double the restart interval each time
+            eta_min=1e-7  # Minimum learning rate
+        )
+        
+        return {
+            'optimizer': optimizer,
+            'lr_scheduler': {
+                'scheduler': scheduler,
+                'monitor': 'val_loss',
+                'interval': 'epoch',
+                'frequency': 1
+            }
+        }
+    
+    def create_callbacks(self) -> List[pl.Callback]:
+        """Create enhanced callbacks."""
+        callbacks = []
+        
+        # Model checkpointing
+        checkpoint_callback = ModelCheckpoint(
+            dirpath='models/checkpoints',
+            filename=f'{self.config.get("experiment_name", "model")}-{{epoch:02d}}-{{val_loss:.4f}}',
+            monitor='val_loss',
+            mode='min',
+            save_top_k=3,
+            save_last=True
+        )
+        callbacks.append(checkpoint_callback)
+        
+        # Learning rate monitoring
+        lr_monitor = LearningRateMonitor(logging_interval='epoch')
+        callbacks.append(lr_monitor)
+        
+        # Early stopping (optional)
+        if self.use_early_stopping:
+            early_stopping = EarlyStopping(
+                monitor='val_loss',
+                mode='min',
+                patience=self.patience,
+                verbose=True
+            )
+            callbacks.append(early_stopping)
+        
+        return callbacks
+    
+    def create_data_loaders(self, data_path: str) -> Tuple[Any, Any, Any]:
+        """Create enhanced data loaders."""
+        # Enhanced data loader with better augmentation
+        data_loader = EnhancedArchitecturalDataLoader(
+            data_dir=data_path,
+            batch_size=self.batch_size,
+            num_workers=4,
+            use_albumentations=True  # Use advanced augmentation
+        )
+        
+        # Calculate sample sizes based on available data
+        total_samples = len(data_loader.get_train_loader().dataset)
+        train_samples = int(0.7 * total_samples)
+        val_samples = max(1, int(0.15 * total_samples))
+        test_samples = max(1, int(0.15 * total_samples))
+        
+        print(f"Data split: Train={train_samples}, Val={val_samples}, Test={test_samples}")
+        
+        train_loader = data_loader.get_train_loader(train_samples)
+        val_loader = data_loader.get_val_loader(val_samples)
+        test_loader = data_loader.get_test_loader(test_samples)
+        
+        return train_loader, val_loader, test_loader
+    
+    def update_curriculum(self, epoch: int):
+        """Update curriculum learning stage."""
+        # Progressive curriculum: start with fewer classes, gradually increase
+        if epoch < 10:
+            self.curriculum_stage = 0
+            self.curriculum_classes_count = min(10, self.num_classes)
+        elif epoch < 30:
+            self.curriculum_stage = 1
+            self.curriculum_classes_count = min(20, self.num_classes)
+        else:
+            self.curriculum_stage = 2
+            self.curriculum_classes_count = self.num_classes
+        
+        # Update model for current curriculum stage
+        self.update_model_for_stage()
+    
+    def update_model_for_stage(self):
+        """Update model for current curriculum stage."""
+        # This can be implemented to modify model behavior based on curriculum stage
+        pass
+
+
+class EnhancedExperimentRunner:
+    """Enhanced experiment runner with advanced optimization."""
+    
+    def __init__(self, config: Dict[str, Any]):
+        self.config = config
+        self.experiment_name = config.get('experiment_name', 'enhanced_experiment')
+        
+    def run_experiment(self, model: nn.Module, data_path: str):
+        """Run enhanced experiment."""
+        print(f"Starting enhanced experiment: {self.experiment_name}")
+        
+        # Create enhanced trainer
+        trainer = EnhancedArchitecturalTrainer(model, self.config)
+        
+        # Create data loaders
+        train_loader, val_loader, test_loader = trainer.create_data_loaders(data_path)
+        
+        # Create callbacks
+        callbacks = trainer.create_callbacks()
+        
+        # Create Lightning trainer
+        lightning_trainer = pl.Trainer(
+            max_epochs=self.config.get('epochs', 100),
+            accelerator='auto',
+            devices='auto',
+            precision='16-mixed' if self.config.get('use_mixed_precision', True) else '32',
+            gradient_clip_val=self.config.get('gradient_clip_val', 1.0),
+            accumulate_grad_batches=self.config.get('accumulate_grad_batches', 2),
+            callbacks=callbacks,
+            logger=trainer.tensorboard_logger,
+            log_every_n_steps=10,
+            val_check_interval=0.5,  # Validate twice per epoch
+            enable_progress_bar=True,
+            enable_model_summary=True,
+            enable_checkpointing=True,
+        )
+        
+        # Train the model
+        lightning_trainer.fit(trainer, train_loader, val_loader)
+        
+        # Test the model
+        lightning_trainer.test(trainer, test_loader)
+        
+        print(f"Enhanced experiment {self.experiment_name} completed successfully!")
+        
+        return trainer
+
+
+# Keep backward compatibility
+class ArchitecturalTrainer(EnhancedArchitecturalTrainer):
+    """Backward compatibility wrapper."""
+    pass
+
+class ExperimentRunner(EnhancedExperimentRunner):
+    """Backward compatibility wrapper."""
+    pass