File size: 4,283 Bytes
6d0498a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
import torch
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
import numpy as np
class GraphMetrics:
"""Production-ready evaluation metrics"""
@staticmethod
def accuracy(pred, target):
"""Classification accuracy"""
pred_labels = pred.argmax(dim=1)
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')
@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()
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')
@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]
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),
}
# 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
@staticmethod
def evaluate_graph_classification(model, dataloader, device):
"""Comprehensive graph classification evaluation"""
model.eval()
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:
num_classes = all_preds.size(1)
metrics['roc_auc'] = GraphMetrics.roc_auc(all_preds, all_targets, num_classes)
except:
pass
return metrics
return {'error': 'No predictions generated'} |