Update core/trainer.py

core/trainer.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torch.optim.lr_scheduler import OneCycleLR, CosineAnnealingWarmRestarts
+from torch.optim.lr_scheduler import ReduceLROnPlateau
 import numpy as np
 import time
 import logging
@@ -10,33 +10,40 @@ from utils.metrics import GraphMetrics
 logger = logging.getLogger(__name__)
 
 class GraphMambaTrainer:
-    """Enhanced trainer with optimized learning rates and schedules"""
+    """Anti-overfitting trainer with heavy regularization"""
     
     def __init__(self, model, config, device):
         self.model = model
         self.config = config
         self.device = device
         
-        # Fixed learning rate (much lower)
-        self.lr = 0.001  # Changed from 0.01
+        # Conservative learning rate
+        self.lr = config['training']['learning_rate']  # Should be 0.0005
         self.epochs = config['training']['epochs']
-        self.patience = config['training'].get('patience', 15)
+        self.patience = config['training'].get('patience', 10)
         self.min_lr = config['training'].get('min_lr', 1e-6)
         
-        # Enhanced optimizer
+        # Heavily regularized optimizer
         self.optimizer = optim.AdamW(
             model.parameters(),
             lr=self.lr,
-            weight_decay=config['training']['weight_decay'],
+            weight_decay=config['training']['weight_decay'],  # Should be 0.01
             betas=(0.9, 0.999),
             eps=1e-8
         )
         
-        # Proper loss function
-        self.criterion = nn.CrossEntropyLoss()
+        # Proper loss function with label smoothing
+        self.criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
         
-        # Learning rate scheduler (will be set in training)
-        self.scheduler = None
+        # Conservative scheduler
+        self.scheduler = ReduceLROnPlateau(
+            self.optimizer,
+            mode='max',
+            factor=0.5,
+            patience=5,
+            min_lr=self.min_lr,
+            verbose=True
+        )
         
         # Training state
         self.best_val_acc = 0.0
@@ -46,36 +53,34 @@ class GraphMambaTrainer:
             'train_loss': [], 'train_acc': [],
             'val_loss': [], 'val_acc': [], 'lr': []
         }
-    
-    def _setup_scheduler(self, total_steps):
-        """Setup learning rate scheduler"""
-        self.scheduler = OneCycleLR(
-            self.optimizer,
-            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
-        )
+        
+        # Track overfitting
+        self.best_gap = float('inf')
+        self.overfitting_threshold = 0.3  # Stop if train-val gap > 30%
     
     def train_node_classification(self, data, verbose=True):
-        """Enhanced training with proper LR scheduling"""
+        """Anti-overfitting training"""
         
         if verbose:
+            total_params = sum(p.numel() for p in self.model.parameters())
+            train_samples = data.train_mask.sum().item()
+            params_per_sample = total_params / train_samples
+            
             print(f"🏋️ Training GraphMamba 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"⚙️ Parameters: {total_params:,}")
+            print(f"📚 Training samples: {train_samples}")
+            print(f"⚠️ Params per sample: {params_per_sample:.1f}")
+            
+            if params_per_sample > 1000:
+                print(f"🚨 WARNING: High params per sample ratio - overfitting risk!")
         
         # 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()
         
@@ -86,6 +91,9 @@ class GraphMambaTrainer:
             # Validation step
             val_metrics = self._validate_epoch(data, epoch)
             
+            # Calculate overfitting gap
+            acc_gap = train_metrics['acc'] - val_metrics['acc']
+            
             # Update history
             self.training_history['train_loss'].append(train_metrics['loss'])
             self.training_history['train_acc'].append(train_metrics['acc'])
@@ -93,59 +101,85 @@ class GraphMambaTrainer:
             self.training_history['val_acc'].append(val_metrics['acc'])
             self.training_history['lr'].append(self.optimizer.param_groups[0]['lr'])
             
+            # Step scheduler
+            self.scheduler.step(val_metrics['acc'])
+            
             # Check for improvement
             if val_metrics['acc'] > self.best_val_acc:
                 self.best_val_acc = val_metrics['acc']
                 self.best_val_loss = val_metrics['loss']
+                self.best_gap = acc_gap
                 self.patience_counter = 0
                 if verbose:
                     print(f"🎉 New best validation accuracy: {self.best_val_acc:.4f}")
             else:
                 self.patience_counter += 1
             
-            # Progress logging
+            # Overfitting detection
+            if acc_gap > self.overfitting_threshold:
+                if verbose:
+                    print(f"🚨 OVERFITTING detected: {acc_gap:.3f} gap")
+                    print(f"   Train: {train_metrics['acc']:.3f}, Val: {val_metrics['acc']:.3f}")
+            
+            # Progress logging with overfitting monitoring
             if verbose and (epoch == 0 or (epoch + 1) % 10 == 0 or epoch == self.epochs - 1):
                 elapsed = time.time() - start_time
+                gap_indicator = "🚨" if acc_gap > 0.2 else "⚠️" if acc_gap > 0.1 else "✅"
+                
                 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")
+                      f"Gap: {acc_gap:.3f} {gap_indicator} | "
+                      f"LR: {self.optimizer.param_groups[0]['lr']:.6f}")
             
-            # Early stopping
+            # Early stopping conditions
             if self.patience_counter >= self.patience:
                 if verbose:
-                    print(f"🛑 Early stopping at epoch {epoch}")
+                    print(f"🛑 Early stopping at epoch {epoch} (patience)")
+                break
+                
+            # Stop if severe overfitting
+            if acc_gap > 0.5:
+                if verbose:
+                    print(f"🛑 Stopping due to severe overfitting (gap: {acc_gap:.3f})")
                 break
-            
-            # Step scheduler
-            self.scheduler.step()
         
         if verbose:
             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}")
+            print(f"📊 Best train-val gap: {self.best_gap:.4f}")
+            
+            if self.best_gap < 0.1:
+                print("🎉 Excellent generalization!")
+            elif self.best_gap < 0.2:
+                print("👍 Good generalization")
+            else:
+                print("⚠️ Some overfitting detected")
         
         return self.training_history
     
     def _train_epoch(self, data, epoch):
-        """Single training epoch"""
+        """Single training epoch with regularization"""
         self.model.train()
         self.optimizer.zero_grad()
         
-        # Forward pass
+        # Forward pass (with data augmentation)
         h = self.model(data.x, data.edge_index)
         logits = self.model.classifier(h)
         
-        # Compute loss on training nodes
+        # Compute loss on training nodes only
         train_loss = self.criterion(logits[data.train_mask], data.y[data.train_mask])
         
-        # Backward pass
-        train_loss.backward()
+        # Add L2 regularization manually if needed
+        l2_reg = 0.0
+        for param in self.model.parameters():
+            l2_reg += torch.norm(param, p=2)
+        train_loss += 1e-5 * l2_reg  # Small additional L2
         
-        # Gradient clipping
+        # Backward pass with gradient clipping
+        train_loss.backward()
         torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
-        
         self.optimizer.step()
         
         # Compute accuracy
@@ -156,7 +190,7 @@ class GraphMambaTrainer:
         return {'loss': train_loss.item(), 'acc': train_acc}
     
     def _validate_epoch(self, data, epoch):
-        """Single validation epoch"""
+        """Validation without augmentation"""
         self.model.eval()
         
         with torch.no_grad():
@@ -171,13 +205,12 @@ class GraphMambaTrainer:
         return {'loss': val_loss.item(), 'acc': val_acc}
     
     def test(self, data):
-        """Comprehensive test evaluation"""
+        """Test evaluation"""
         self.model.eval()
         
         with torch.no_grad():
             h = self.model(data.x, data.edge_index)
             
-            # Ensure classifier exists
             if self.model.classifier is None:
                 num_classes = len(torch.unique(data.y))
                 self.model._init_classifier(num_classes, self.device)
@@ -189,7 +222,6 @@ class GraphMambaTrainer:
             test_pred = logits[data.test_mask]
             test_target = data.y[data.test_mask]
             
-            # Comprehensive metrics
             metrics = {
                 'test_loss': test_loss.item(),
                 'test_acc': GraphMetrics.accuracy(test_pred, test_target),
@@ -197,7 +229,6 @@ class GraphMambaTrainer:
                 'f1_micro': GraphMetrics.f1_score_micro(test_pred, test_target),
             }
             
-            # Additional metrics
             precision, recall = GraphMetrics.precision_recall(test_pred, test_target)
             metrics['precision'] = precision
             metrics['recall'] = recall
@@ -208,64 +239,4 @@ class GraphMambaTrainer:
         """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)
-        
-        # 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 train_node_classification(self, data, verbose=True):
-        """Training with enhanced monitoring"""
-        
-        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)")
-        
-        # Call parent method with enhancements
-        history = super().train_node_classification(data, verbose)
-        
-        # 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}%)")
-            
-            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 history
+            return self.model(data.x, data.edge_index)