Update utils/metrics.py

utils/metrics.py

@@ -4,73 +4,84 @@ from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
 import numpy as np
 
 class GraphMetrics:
-    """Production-ready evaluation metrics"""
+    """Production-ready evaluation metrics - device safe"""
     
     @staticmethod
     def accuracy(pred, target):
         """Classification accuracy"""
-        pred_labels = pred.argmax(dim=1)
+        if pred.dim() > 1:
+            pred_labels = pred.argmax(dim=1)
+        else:
+            pred_labels = pred
         return (pred_labels == target).float().mean().item()
     
     @staticmethod 
     def f1_score_macro(pred, target):
         """Macro F1 score"""
-        pred_labels = pred.argmax(dim=1).cpu().numpy()
-        target_labels = target.cpu().numpy()
-        return f1_score(target_labels, pred_labels, average='macro')
+        try:
+            if pred.dim() > 1:
+                pred_labels = pred.argmax(dim=1)
+            else:
+                pred_labels = pred
+            pred_labels = pred_labels.cpu().numpy()
+            target_labels = target.cpu().numpy()
+            return f1_score(target_labels, pred_labels, average='macro', zero_division=0)
+        except:
+            return 0.0
     
     @staticmethod
     def f1_score_micro(pred, target):
         """Micro F1 score"""
-        pred_labels = pred.argmax(dim=1).cpu().numpy() 
-        target_labels = target.cpu().numpy()
-        return f1_score(target_labels, pred_labels, average='micro')
-    
-    @staticmethod
-    def roc_auc(pred, target, num_classes):
-        """ROC AUC for multi-class"""
-        if num_classes == 2:
-            # Binary classification
-            pred_probs = F.softmax(pred, dim=1)[:, 1].cpu().numpy()
+        try:
+            if pred.dim() > 1:
+                pred_labels = pred.argmax(dim=1)
+            else:
+                pred_labels = pred
+            pred_labels = pred_labels.cpu().numpy() 
             target_labels = target.cpu().numpy()
-            return roc_auc_score(target_labels, pred_probs)
-        else:
-            # Multi-class
-            pred_probs = F.softmax(pred, dim=1).cpu().numpy()
-            target_onehot = F.one_hot(target, num_classes).cpu().numpy()
-            return roc_auc_score(target_onehot, pred_probs, multi_class='ovr', average='macro')
+            return f1_score(target_labels, pred_labels, average='micro', zero_division=0)
+        except:
+            return 0.0
     
     @staticmethod
     def evaluate_node_classification(model, data, mask, device):
         """Comprehensive node classification evaluation"""
         model.eval()
         
-        with torch.no_grad():
-            data = data.to(device)
-            h = model(data.x, data.edge_index)
-            
-            # Assuming a classification head exists
-            if hasattr(model, 'classifier'):
-                pred = model.classifier(h)
-            else:
-                # If no classifier, return embeddings
-                return {'embeddings': h[mask].cpu()}
-            
-            pred_masked = pred[mask]
-            target_masked = data.y[mask]
-            
+        try:
+            with torch.no_grad():
+                # Ensure data is on correct device
+                data = data.to(device)
+                model = model.to(device)
+                
+                h = model(data.x, data.edge_index)
+                
+                # Get predictions
+                if hasattr(model, 'classifier') and model.classifier is not None:
+                    pred = model.classifier(h)
+                else:
+                    # Initialize classifier if needed
+                    num_classes = len(torch.unique(data.y))
+                    model._init_classifier(num_classes, device)
+                    pred = model.classifier(h)
+                
+                pred_masked = pred[mask]
+                target_masked = data.y[mask]
+                
+                metrics = {
+                    'accuracy': GraphMetrics.accuracy(pred_masked, target_masked),
+                    'f1_macro': GraphMetrics.f1_score_macro(pred_masked, target_masked),
+                    'f1_micro': GraphMetrics.f1_score_micro(pred_masked, target_masked),
+                }
+                
+        except Exception as e:
+            print(f"Evaluation error: {e}")
             metrics = {
-                'accuracy': GraphMetrics.accuracy(pred_masked, target_masked),
-                'f1_macro': GraphMetrics.f1_score_macro(pred_masked, target_masked),
-                'f1_micro': GraphMetrics.f1_score_micro(pred_masked, target_masked),
+                'accuracy': 0.0,
+                'f1_macro': 0.0,
+                'f1_micro': 0.0,
+                'error': str(e)
             }
-            
-            # Add ROC AUC if binary/multi-class
-            try:
-                num_classes = pred.size(1)
-                metrics['roc_auc'] = GraphMetrics.roc_auc(pred_masked, target_masked, num_classes)
-            except:
-                pass
                 
         return metrics
     
@@ -82,35 +93,40 @@ class GraphMetrics:
         all_preds = []
         all_targets = []
         
-        with torch.no_grad():
-            for batch in dataloader:
-                batch = batch.to(device)
-                h = model(batch.x, batch.edge_index, batch.batch)
-                
-                # Graph-level prediction
-                graph_h = model.get_graph_embedding(h, batch.batch)
-                
-                if hasattr(model, 'classifier'):
-                    pred = model.classifier(graph_h)
-                    all_preds.append(pred)
-                    all_targets.append(batch.y)
-        
-        if all_preds:
-            all_preds = torch.cat(all_preds, dim=0)
-            all_targets = torch.cat(all_targets, dim=0)
-            
-            metrics = {
-                'accuracy': GraphMetrics.accuracy(all_preds, all_targets),
-                'f1_macro': GraphMetrics.f1_score_macro(all_preds, all_targets),
-                'f1_micro': GraphMetrics.f1_score_micro(all_preds, all_targets),
-            }
+        try:
+            with torch.no_grad():
+                for batch in dataloader:
+                    batch = batch.to(device)
+                    h = model(batch.x, batch.edge_index, batch.batch)
+                    
+                    # Graph-level prediction
+                    graph_h = model.get_graph_embedding(h, batch.batch)
+                    
+                    if hasattr(model, 'classifier') and model.classifier is not None:
+                        pred = model.classifier(graph_h)
+                    else:
+                        # Initialize classifier
+                        num_classes = len(torch.unique(batch.y))
+                        model._init_classifier(num_classes, device)
+                        pred = model.classifier(graph_h)
+                    
+                    all_preds.append(pred.cpu())
+                    all_targets.append(batch.y.cpu())
             
-            try:
-                num_classes = all_preds.size(1)
-                metrics['roc_auc'] = GraphMetrics.roc_auc(all_preds, all_targets, num_classes)
-            except:
-                pass
+            if all_preds:
+                all_preds = torch.cat(all_preds, dim=0)
+                all_targets = torch.cat(all_targets, dim=0)
+                
+                metrics = {
+                    'accuracy': GraphMetrics.accuracy(all_preds, all_targets),
+                    'f1_macro': GraphMetrics.f1_score_macro(all_preds, all_targets),
+                    'f1_micro': GraphMetrics.f1_score_micro(all_preds, all_targets),
+                }
+            else:
+                metrics = {'error': 'No predictions generated'}
                 
-            return metrics
+        except Exception as e:
+            print(f"Graph classification evaluation error: {e}")
+            metrics = {'error': str(e)}
         
-        return {'error': 'No predictions generated'}
+        return metrics