Update core/trainer.py

core/trainer.py

@@ -18,7 +18,7 @@ class GraphMambaTrainer:
         self.device = device
         
         # Conservative learning rate
-        self.lr = config['training']['learning_rate']  # Should be 0.0005
+        self.lr = config['training']['learning_rate']
         self.epochs = config['training']['epochs']
         self.patience = config['training'].get('patience', 10)
         self.min_lr = config['training'].get('min_lr', 1e-6)
@@ -27,7 +27,7 @@ class GraphMambaTrainer:
         self.optimizer = optim.AdamW(
             model.parameters(),
             lr=self.lr,
-            weight_decay=config['training']['weight_decay'],  # Should be 0.01
+            weight_decay=config['training']['weight_decay'],
             betas=(0.9, 0.999),
             eps=1e-8
         )
@@ -35,14 +35,13 @@ class GraphMambaTrainer:
         # Proper loss function with label smoothing
         self.criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
         
-        # Conservative scheduler
+        # Conservative scheduler - FIXED: removed verbose parameter
         self.scheduler = ReduceLROnPlateau(
             self.optimizer,
             mode='max',
             factor=0.5,
             patience=5,
-            min_lr=self.min_lr,
-            verbose=True
+            min_lr=self.min_lr
         )
         
         # Training state
@@ -56,7 +55,7 @@ class GraphMambaTrainer:
         
         # Track overfitting
         self.best_gap = float('inf')
-        self.overfitting_threshold = 0.3  # Stop if train-val gap > 30%
+        self.overfitting_threshold = 0.3
     
     def train_node_classification(self, data, verbose=True):
         """Anti-overfitting training"""
@@ -121,7 +120,7 @@ class GraphMambaTrainer:
                     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
+            # Progress logging
             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 "✅"
@@ -171,11 +170,11 @@ class GraphMambaTrainer:
         # Compute loss on training nodes only
         train_loss = self.criterion(logits[data.train_mask], data.y[data.train_mask])
         
-        # Add L2 regularization manually if needed
+        # Add L2 regularization manually
         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
+        train_loss += 1e-5 * l2_reg
         
         # Backward pass with gradient clipping
         train_loss.backward()