Update core/trainer.py

core/trainer.py

@@ -1,299 +1,271 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau
+from torch.optim.lr_scheduler import OneCycleLR, CosineAnnealingWarmRestarts
 import numpy as np
-from tqdm import tqdm
 import time
-import matplotlib.pyplot as plt
+import logging
+from utils.metrics import GraphMetrics
+
+logger = logging.getLogger(__name__)
 
 class GraphMambaTrainer:
-    """
-    Production-ready trainer for GraphMamba
-    Includes advanced training techniques
-    """
+    """Enhanced trainer with optimized learning rates and schedules"""
     
-    def __init__(self, model, config, device='cpu'):
-        self.model = model.to(device)
+    def __init__(self, model, config, device):
+        self.model = model
         self.config = config
         self.device = device
         
-        # Training parameters
-        self.lr = config['training']['learning_rate']
-        self.weight_decay = config['training']['weight_decay']
+        # Fixed learning rate (much lower)
+        self.lr = 0.001  # Changed from 0.01
         self.epochs = config['training']['epochs']
-        self.patience = config['training']['patience']
-        self.warmup_epochs = config['training']['warmup_epochs']
-        self.min_lr = config['training']['min_lr']
+        self.patience = config['training'].get('patience', 15)
+        self.min_lr = config['training'].get('min_lr', 1e-6)
         
-        # Initialize optimizer
+        # Enhanced optimizer
         self.optimizer = optim.AdamW(
-            self.model.parameters(),
+            model.parameters(),
             lr=self.lr,
-            weight_decay=self.weight_decay,
+            weight_decay=config['training']['weight_decay'],
             betas=(0.9, 0.999),
             eps=1e-8
         )
         
-        # Loss function
+        # Proper loss function
         self.criterion = nn.CrossEntropyLoss()
         
-        # Scheduler
+        # Learning rate scheduler (will be set in training)
         self.scheduler = None
         
-        # Training history
-        self.history = {
-            'train_loss': [],
-            'train_acc': [],
-            'val_loss': [],
-            'val_acc': [],
-            'lr': []
-        }
-        
-        # Best model tracking
+        # Training state
         self.best_val_acc = 0.0
-        self.best_model_state = None
+        self.best_val_loss = float('inf')
         self.patience_counter = 0
-        
-    def train_node_classification(self, data, verbose=True):
-        """
-        Train model for node classification
-        """
-        # Initialize classifier
-        num_classes = len(torch.unique(data.y))
-        self.model._init_classifier(num_classes, self.device)
-        
-        # Update optimizer to include new parameters
-        self.optimizer = optim.AdamW(
-            self.model.parameters(),
-            lr=self.lr,
-            weight_decay=self.weight_decay,
-            betas=(0.9, 0.999)
-        )
-        
-        # Initialize scheduler
-        self.scheduler = CosineAnnealingLR(
+        self.training_history = {
+            'train_loss': [], 'train_acc': [],
+            'val_loss': [], 'val_acc': [], 'lr': []
+        }
+    
+    def _setup_scheduler(self, total_steps):
+        """Setup learning rate scheduler"""
+        self.scheduler = OneCycleLR(
             self.optimizer,
-            T_max=self.epochs - self.warmup_epochs,
-            eta_min=self.min_lr
+            max_lr=self.lr,
+            total_steps=total_steps,
+            pct_start=0.1,  # 10% warmup
+            anneal_strategy='cos',
+            div_factor=10.0,  # Start LR = max_lr/10
+            final_div_factor=100.0  # End LR = max_lr/100
         )
+    
+    def train_node_classification(self, data, verbose=True):
+        """Enhanced training with proper LR scheduling"""
         
         if verbose:
             print(f"🏋️ Training GraphMamba for {self.epochs} epochs")
             print(f"📊 Dataset: {data.num_nodes} nodes, {data.num_edges} edges")
-            print(f"🎯 Classes: {num_classes}")
+            print(f"🎯 Classes: {len(torch.unique(data.y))}")
             print(f"💾 Device: {self.device}")
             print(f"⚙️ Parameters: {sum(p.numel() for p in self.model.parameters()):,}")
         
-        # Training loop
+        # Initialize classifier
+        num_classes = len(torch.unique(data.y))
+        self.model._init_classifier(num_classes, self.device)
+        
+        # Setup scheduler
+        self._setup_scheduler(self.epochs)
+        
+        self.model.train()
+        start_time = time.time()
+        
         for epoch in range(self.epochs):
-            # Training phase
-            train_loss, train_acc = self._train_epoch(data, epoch)
+            # Training step
+            train_metrics = self._train_epoch(data, epoch)
             
-            # Validation phase
-            val_loss, val_acc = self._validate_epoch(data)
+            # Validation step
+            val_metrics = self._validate_epoch(data, epoch)
             
-            # Learning rate scheduling
-            if epoch >= self.warmup_epochs:
-                self.scheduler.step()
-            else:
-                # Warmup
-                warmup_lr = self.lr * (epoch + 1) / self.warmup_epochs
-                for param_group in self.optimizer.param_groups:
-                    param_group['lr'] = warmup_lr
-            
-            # Record history
-            current_lr = self.optimizer.param_groups[0]['lr']
-            self.history['train_loss'].append(train_loss)
-            self.history['train_acc'].append(train_acc)
-            self.history['val_loss'].append(val_loss)
-            self.history['val_acc'].append(val_acc)
-            self.history['lr'].append(current_lr)
+            # Update history
+            self.training_history['train_loss'].append(train_metrics['loss'])
+            self.training_history['train_acc'].append(train_metrics['acc'])
+            self.training_history['val_loss'].append(val_metrics['loss'])
+            self.training_history['val_acc'].append(val_metrics['acc'])
+            self.training_history['lr'].append(self.optimizer.param_groups[0]['lr'])
             
             # Check for improvement
-            if val_acc > self.best_val_acc:
-                self.best_val_acc = val_acc
-                self.best_model_state = self.model.state_dict().copy()
+            if val_metrics['acc'] > self.best_val_acc:
+                self.best_val_acc = val_metrics['acc']
+                self.best_val_loss = val_metrics['loss']
                 self.patience_counter = 0
-                
-                if verbose and epoch % 10 == 0:
-                    print(f"🎉 New best validation accuracy: {val_acc:.4f}")
+                if verbose:
+                    print(f"🎉 New best validation accuracy: {self.best_val_acc:.4f}")
             else:
                 self.patience_counter += 1
             
+            # Progress logging
+            if verbose and (epoch == 0 or (epoch + 1) % 10 == 0 or epoch == self.epochs - 1):
+                elapsed = time.time() - start_time
+                print(f"Epoch {epoch:3d} | "
+                      f"Train: {train_metrics['loss']:.4f} ({train_metrics['acc']:.4f}) | "
+                      f"Val: {val_metrics['loss']:.4f} ({val_metrics['acc']:.4f}) | "
+                      f"LR: {self.optimizer.param_groups[0]['lr']:.6f} | "
+                      f"Time: {elapsed:.1f}s")
+            
             # Early stopping
             if self.patience_counter >= self.patience:
                 if verbose:
-                    print(f"⏹️ Early stopping at epoch {epoch}")
+                    print(f"🛑 Early stopping at epoch {epoch}")
                 break
             
-            # Progress reporting
-            if verbose and epoch % 20 == 0:
-                print(f"Epoch {epoch:3d} | "
-                      f"Train: {train_loss:.4f} ({train_acc:.4f}) | "
-                      f"Val: {val_loss:.4f} ({val_acc:.4f}) | "
-                      f"LR: {current_lr:.6f}")
+            # Step scheduler
+            self.scheduler.step()
         
-        # Load best model
-        if self.best_model_state is not None:
-            self.model.load_state_dict(self.best_model_state)
-            
         if verbose:
-            print(f"✅ Training completed!")
+            total_time = time.time() - start_time
+            print(f"✅ Training completed in {total_time:.2f}s")
             print(f"🏆 Best validation accuracy: {self.best_val_acc:.4f}")
         
-        return self.history
+        return self.training_history
     
     def _train_epoch(self, data, epoch):
         """Single training epoch"""
         self.model.train()
-        
-        # Forward pass
         self.optimizer.zero_grad()
         
+        # Forward pass
         h = self.model(data.x, data.edge_index)
-        pred = self.model.classifier(h)
+        logits = self.model.classifier(h)
         
-        # Loss only on training nodes
-        loss = self.criterion(pred[data.train_mask], data.y[data.train_mask])
+        # Compute loss on training nodes
+        train_loss = self.criterion(logits[data.train_mask], data.y[data.train_mask])
         
         # Backward pass
-        loss.backward()
+        train_loss.backward()
         
         # Gradient clipping
         torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
         
         self.optimizer.step()
         
-        # Calculate accuracy
+        # Compute accuracy
         with torch.no_grad():
-            pred_labels = pred[data.train_mask].argmax(dim=1)
-            acc = (pred_labels == data.y[data.train_mask]).float().mean()
+            train_pred = logits[data.train_mask].argmax(dim=1)
+            train_acc = (train_pred == data.y[data.train_mask]).float().mean().item()
         
-        return loss.item(), acc.item()
+        return {'loss': train_loss.item(), 'acc': train_acc}
     
-    def _validate_epoch(self, data):
+    def _validate_epoch(self, data, epoch):
         """Single validation epoch"""
         self.model.eval()
         
         with torch.no_grad():
             h = self.model(data.x, data.edge_index)
-            pred = self.model.classifier(h)
-            
-            # Loss on validation nodes
-            val_loss = self.criterion(pred[data.val_mask], data.y[data.val_mask])
+            logits = self.model.classifier(h)
             
-            # Accuracy
-            pred_labels = pred[data.val_mask].argmax(dim=1)
-            val_acc = (pred_labels == data.y[data.val_mask]).float().mean()
+            # Validation loss and accuracy
+            val_loss = self.criterion(logits[data.val_mask], data.y[data.val_mask])
+            val_pred = logits[data.val_mask].argmax(dim=1)
+            val_acc = (val_pred == data.y[data.val_mask]).float().mean().item()
         
-        return val_loss.item(), val_acc.item()
+        return {'loss': val_loss.item(), 'acc': val_acc}
     
     def test(self, data):
-        """Test the model"""
+        """Comprehensive test evaluation"""
         self.model.eval()
         
         with torch.no_grad():
             h = self.model(data.x, data.edge_index)
-            pred = self.model.classifier(h)
+            
+            # Ensure classifier exists
+            if self.model.classifier is None:
+                num_classes = len(torch.unique(data.y))
+                self.model._init_classifier(num_classes, self.device)
+            
+            logits = self.model.classifier(h)
             
             # Test metrics
-            test_loss = self.criterion(pred[data.test_mask], data.y[data.test_mask])
-            pred_labels = pred[data.test_mask].argmax(dim=1)
-            test_acc = (pred_labels == data.y[data.test_mask]).float().mean()
+            test_loss = self.criterion(logits[data.test_mask], data.y[data.test_mask])
+            test_pred = logits[data.test_mask]
+            test_target = data.y[data.test_mask]
             
-            # Per-class accuracy
-            num_classes = len(torch.unique(data.y))
-            class_acc = []
+            # Comprehensive metrics
+            metrics = {
+                'test_loss': test_loss.item(),
+                'test_acc': GraphMetrics.accuracy(test_pred, test_target),
+                'f1_macro': GraphMetrics.f1_score_macro(test_pred, test_target),
+                'f1_micro': GraphMetrics.f1_score_micro(test_pred, test_target),
+            }
             
-            for c in range(num_classes):
-                class_mask = data.y[data.test_mask] == c
-                if class_mask.any():
-                    class_correct = (pred_labels[class_mask] == c).float().mean()
-                    class_acc.append(class_correct.item())
-                else:
-                    class_acc.append(0.0)
-        
-        return {
-            'test_loss': test_loss.item(),
-            'test_acc': test_acc.item(),
-            'class_acc': class_acc
-        }
+            # Additional metrics
+            precision, recall = GraphMetrics.precision_recall(test_pred, test_target)
+            metrics['precision'] = precision
+            metrics['recall'] = recall
+            
+        return metrics
     
-    def plot_training_history(self, save_path=None):
-        """Plot training history"""
-        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(12, 8))
-        
-        epochs = range(len(self.history['train_loss']))
-        
-        # Loss plot
-        ax1.plot(epochs, self.history['train_loss'], label='Train', color='blue')
-        ax1.plot(epochs, self.history['val_loss'], label='Validation', color='red')
-        ax1.set_title('Training Loss')
-        ax1.set_xlabel('Epoch')
-        ax1.set_ylabel('Loss')
-        ax1.legend()
-        ax1.grid(True, alpha=0.3)
-        
-        # Accuracy plot
-        ax2.plot(epochs, self.history['train_acc'], label='Train', color='blue')
-        ax2.plot(epochs, self.history['val_acc'], label='Validation', color='red')
-        ax2.set_title('Training Accuracy')
-        ax2.set_xlabel('Epoch')
-        ax2.set_ylabel('Accuracy')
-        ax2.legend()
-        ax2.grid(True, alpha=0.3)
-        
-        # Learning rate plot
-        ax3.plot(epochs, self.history['lr'], color='green')
-        ax3.set_title('Learning Rate')
-        ax3.set_xlabel('Epoch')
-        ax3.set_ylabel('Learning Rate')
-        ax3.set_yscale('log')
-        ax3.grid(True, alpha=0.3)
-        
-        # Best metrics
-        best_train_acc = max(self.history['train_acc'])
-        best_val_acc = max(self.history['val_acc'])
-        
-        ax4.bar(['Best Train Acc', 'Best Val Acc'], [best_train_acc, best_val_acc], 
-                color=['blue', 'red'], alpha=0.7)
-        ax4.set_title('Best Accuracies')
-        ax4.set_ylabel('Accuracy')
-        ax4.set_ylim(0, 1)
-        
-        for i, v in enumerate([best_train_acc, best_val_acc]):
-            ax4.text(i, v + 0.01, f'{v:.4f}', ha='center', va='bottom')
-        
-        plt.tight_layout()
-        
-        if save_path:
-            plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    def get_embeddings(self, data):
+        """Get node embeddings"""
+        self.model.eval()
+        with torch.no_grad():
+            return self.model(data.x, data.edge_index)
+
+
+class EnhancedGraphMambaTrainer(GraphMambaTrainer):
+    """Enhanced trainer with additional optimizations"""
+    
+    def __init__(self, model, config, device):
+        super().__init__(model, config, device)
         
-        return fig
+        # Even more conservative learning rate for complex architectures
+        if hasattr(model, 'multi_scale') or 'Hybrid' in model.__class__.__name__:
+            self.lr = 0.0005  # Lower for complex models
+            
+            self.optimizer = optim.AdamW(
+                model.parameters(),
+                lr=self.lr,
+                weight_decay=config['training']['weight_decay'],
+                betas=(0.9, 0.99),  # More stable
+                eps=1e-8
+            )
+    
+    def _setup_scheduler(self, total_steps):
+        """Enhanced scheduler for complex models"""
+        # Cosine annealing with warm restarts
+        self.scheduler = CosineAnnealingWarmRestarts(
+            self.optimizer,
+            T_0=20,  # Restart every 20 epochs
+            T_mult=2,  # Double period after restart
+            eta_min=self.min_lr
+        )
     
-    def save_model(self, path):
-        """Save model and training state"""
-        torch.save({
-            '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,
-            'best_val_acc': self.best_val_acc,
-            'history': self.history,
-            'config': self.config
-        }, path)
+    def train_node_classification(self, data, verbose=True):
+        """Training with enhanced monitoring"""
         
-    def load_model(self, path):
-        """Load model and training state"""
-        checkpoint = torch.load(path, map_location=self.device)
+        if verbose:
+            model_type = self.model.__class__.__name__
+            print(f"🏋️ Training {model_type} for {self.epochs} epochs")
+            print(f"📊 Dataset: {data.num_nodes} nodes, {data.num_edges} edges")
+            print(f"🎯 Classes: {len(torch.unique(data.y))}")
+            print(f"💾 Device: {self.device}")
+            print(f"⚙️ Parameters: {sum(p.numel() for p in self.model.parameters()):,}")
+            print(f"📈 Learning Rate: {self.lr} (enhanced schedule)")
         
-        self.model.load_state_dict(checkpoint['model_state_dict'])
-        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        # Call parent method with enhancements
+        history = super().train_node_classification(data, verbose)
         
-        if checkpoint['scheduler_state_dict'] and self.scheduler:
-            self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+        # Additional analysis
+        if verbose:
+            final_acc = history['val_acc'][-1] if history['val_acc'] else 0
+            improvement = final_acc - (history['val_acc'][0] if history['val_acc'] else 0)
+            print(f"📊 Final validation accuracy: {final_acc:.4f}")
+            print(f"📈 Total improvement: {improvement:.4f} ({improvement*100:.1f}%)")
             
-        self.best_val_acc = checkpoint['best_val_acc']
-        self.history = checkpoint['history']
+            if final_acc > 0.6:
+                print("🎉 Excellent performance! Model converged well.")
+            elif final_acc > 0.4:
+                print("👍 Good progress! Consider more epochs or tuning.")
+            else:
+                print("⚠️ Low accuracy. Check model architecture or data.")
         
-        return checkpoint['config']
+        return history