aibt-fl 1.0.0__py3-none-any.whl

aibt/metrics.py

@@ -0,0 +1,383 @@
+"""
+Privacy Attack Metrics for AIBT Framework
+
+Implements Membership Inference and Attribute Inference attacks for
+evaluating privacy protection of federated learning models.
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+from typing import Dict, Tuple, Optional, List
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import roc_auc_score, accuracy_score, precision_recall_fscore_support
+import warnings
+warnings.filterwarnings('ignore')
+
+
+# ============================================================================
+# MEMBERSHIP INFERENCE ATTACK
+# ============================================================================
+
+class MembershipInferenceAttack:
+    """
+    Membership Inference Attack (MIA) for evaluating privacy leakage.
+    
+    Reference: Shokri et al., "Membership Inference Attacks Against Machine Learning Models"
+    
+    The attack attempts to determine whether a given sample was part of the
+    training data, using model confidence and prediction patterns.
+    """
+    
+    def __init__(self, target_model: nn.Module, device: str = "cpu"):
+        self.target_model = target_model
+        self.device = device
+        self.attack_model = None
+    
+    def get_prediction_features(
+        self,
+        model: nn.Module,
+        X: np.ndarray,
+        y: np.ndarray
+    ) -> np.ndarray:
+        """
+        Extract features from model predictions for attack model.
+        
+        Features include:
+        - Prediction probabilities
+        - Prediction confidence (max prob)
+        - Prediction entropy
+        - Whether prediction is correct
+        - Loss value
+        """
+        model.eval()
+        X_tensor = torch.FloatTensor(X).to(self.device)
+        y_tensor = torch.LongTensor(y).to(self.device)
+        
+        with torch.no_grad():
+            output = model(X_tensor)
+            if isinstance(output, tuple):
+                logits = output[0]
+            else:
+                logits = output
+            
+            probs = F.softmax(logits, dim=1)
+            
+            # Features
+            confidence = probs.max(dim=1)[0].cpu().numpy()
+            correct = (probs.argmax(dim=1) == y_tensor).float().cpu().numpy()
+            entropy = -(probs * torch.log(probs + 1e-10)).sum(dim=1).cpu().numpy()
+            
+            # Per-sample loss
+            loss_per_sample = F.cross_entropy(logits, y_tensor, reduction='none')
+            loss = loss_per_sample.cpu().numpy()
+        
+        # Combine features
+        features = np.column_stack([confidence, correct, entropy, loss])
+        return features
+    
+    def train_attack_model(
+        self,
+        train_data: Tuple[np.ndarray, np.ndarray],
+        test_data: Tuple[np.ndarray, np.ndarray],
+        shadow_model: Optional[nn.Module] = None
+    ) -> float:
+        """
+        Train the attack model using shadow model technique.
+        
+        Args:
+            train_data: (X, y) used for training target model (members)
+            test_data: (X, y) not used for training (non-members)
+            shadow_model: Optional shadow model (uses target if not provided)
+        
+        Returns:
+            AUC score of attack model
+        """
+        model = shadow_model if shadow_model else self.target_model
+        
+        # Get features for members and non-members
+        X_train, y_train = train_data
+        X_test, y_test = test_data
+        
+        member_features = self.get_prediction_features(model, X_train, y_train)
+        non_member_features = self.get_prediction_features(model, X_test, y_test)
+        
+        # Create attack dataset
+        X_attack = np.vstack([member_features, non_member_features])
+        y_attack = np.concatenate([
+            np.ones(len(member_features)),
+            np.zeros(len(non_member_features))
+        ])
+        
+        # Train attack model
+        X_attack_train, X_attack_test, y_attack_train, y_attack_test = train_test_split(
+            X_attack, y_attack, test_size=0.3, random_state=42, stratify=y_attack
+        )
+        
+        self.attack_model = LogisticRegression(max_iter=1000, random_state=42)
+        self.attack_model.fit(X_attack_train, y_attack_train)
+        
+        # Evaluate
+        y_pred_proba = self.attack_model.predict_proba(X_attack_test)[:, 1]
+        auc = roc_auc_score(y_attack_test, y_pred_proba)
+        
+        return auc
+    
+    def attack(
+        self,
+        X: np.ndarray,
+        y: np.ndarray
+    ) -> Tuple[np.ndarray, float]:
+        """
+        Perform membership inference attack.
+        
+        Returns:
+            membership_probs: Probability of being a member
+            auc: Area under ROC curve
+        """
+        if self.attack_model is None:
+            raise ValueError("Attack model not trained. Call train_attack_model first.")
+        
+        features = self.get_prediction_features(self.target_model, X, y)
+        membership_probs = self.attack_model.predict_proba(features)[:, 1]
+        
+        return membership_probs, 0.5  # AUC requires ground truth
+
+
+def evaluate_membership_inference(
+    model: nn.Module,
+    train_data: Tuple[np.ndarray, np.ndarray],
+    test_data: Tuple[np.ndarray, np.ndarray],
+    device: str = "cpu"
+) -> Dict[str, float]:
+    """
+    Complete membership inference evaluation.
+    
+    Returns:
+        Dictionary with AUC and other metrics
+    """
+    attack = MembershipInferenceAttack(model, device)
+    auc = attack.train_attack_model(train_data, test_data)
+    
+    return {
+        "membership_auc": auc,
+        "privacy_score": 1.0 - abs(auc - 0.5) * 2  # Closer to 0.5 = better privacy
+    }
+
+
+# ============================================================================
+# ATTRIBUTE INFERENCE ATTACK
+# ============================================================================
+
+class AttributeInferenceAttack:
+    """
+    Attribute Inference Attack for evaluating sensitive attribute leakage.
+    
+    Attempts to predict sensitive attributes from latent representations.
+    """
+    
+    def __init__(self, target_model: nn.Module, device: str = "cpu"):
+        self.target_model = target_model
+        self.device = device
+        self.attack_model = None
+    
+    def get_latent_representations(
+        self,
+        X: np.ndarray
+    ) -> np.ndarray:
+        """Extract latent representations from the model"""
+        self.target_model.eval()
+        X_tensor = torch.FloatTensor(X).to(self.device)
+        
+        with torch.no_grad():
+            if hasattr(self.target_model, 'get_latent'):
+                z = self.target_model.get_latent(X_tensor)
+            elif hasattr(self.target_model, 'encoder'):
+                encoder_out = self.target_model.encoder(X_tensor)
+                if isinstance(encoder_out, tuple):
+                    z = encoder_out[0]
+                else:
+                    z = encoder_out
+            else:
+                # Use output directly
+                output = self.target_model(X_tensor)
+                if isinstance(output, tuple):
+                    z = output[0]
+                else:
+                    z = output
+        
+        return z.cpu().numpy()
+    
+    def train_attack_model(
+        self,
+        X: np.ndarray,
+        sensitive_attrs: np.ndarray
+    ) -> float:
+        """
+        Train attack model to predict sensitive attributes from latent codes.
+        
+        Args:
+            X: Input features
+            sensitive_attrs: Sensitive attribute labels
+        
+        Returns:
+            AUC score of attribute inference
+        """
+        # Get latent representations
+        latent = self.get_latent_representations(X)
+        
+        # Train-test split
+        z_train, z_test, s_train, s_test = train_test_split(
+            latent, sensitive_attrs, test_size=0.3, random_state=42,
+            stratify=sensitive_attrs
+        )
+        
+        # Train attack model
+        self.attack_model = LogisticRegression(max_iter=1000, random_state=42)
+        self.attack_model.fit(z_train, s_train)
+        
+        # Evaluate
+        s_pred_proba = self.attack_model.predict_proba(z_test)
+        
+        # Handle binary and multi-class
+        if len(np.unique(sensitive_attrs)) == 2:
+            auc = roc_auc_score(s_test, s_pred_proba[:, 1])
+        else:
+            auc = roc_auc_score(s_test, s_pred_proba, multi_class='ovr', average='weighted')
+        
+        return auc
+    
+    def attack(self, X: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Perform attribute inference attack.
+        
+        Returns:
+            predictions: Predicted sensitive attributes
+            probabilities: Prediction probabilities
+        """
+        if self.attack_model is None:
+            raise ValueError("Attack model not trained. Call train_attack_model first.")
+        
+        latent = self.get_latent_representations(X)
+        predictions = self.attack_model.predict(latent)
+        probabilities = self.attack_model.predict_proba(latent)
+        
+        return predictions, probabilities
+
+
+def evaluate_attribute_inference(
+    model: nn.Module,
+    X: np.ndarray,
+    sensitive_attrs: np.ndarray,
+    device: str = "cpu"
+) -> Dict[str, float]:
+    """
+    Complete attribute inference evaluation.
+    
+    Returns:
+        Dictionary with AUC and accuracy
+    """
+    attack = AttributeInferenceAttack(model, device)
+    auc = attack.train_attack_model(X, sensitive_attrs)
+    
+    predictions, _ = attack.attack(X)
+    accuracy = accuracy_score(sensitive_attrs, predictions)
+    
+    return {
+        "attribute_auc": auc,
+        "attribute_accuracy": accuracy,
+        "privacy_score": 1.0 - abs(auc - 0.5) * 2
+    }
+
+
+# ============================================================================
+# COMPREHENSIVE PRIVACY EVALUATION
+# ============================================================================
+
+def evaluate_privacy(
+    model: nn.Module,
+    train_data: Tuple[np.ndarray, np.ndarray],
+    test_data: Tuple[np.ndarray, np.ndarray],
+    sensitive_attrs_train: Optional[np.ndarray] = None,
+    sensitive_attrs_test: Optional[np.ndarray] = None,
+    device: str = "cpu"
+) -> Dict[str, float]:
+    """
+    Comprehensive privacy evaluation.
+    
+    Performs:
+    1. Membership Inference Attack
+    2. Attribute Inference Attack (if sensitive attrs provided)
+    
+    Returns:
+        Dictionary with all privacy metrics
+    """
+    results = {}
+    
+    # Membership Inference Attack
+    mia_metrics = evaluate_membership_inference(model, train_data, test_data, device)
+    results.update({
+        "membership_auc": mia_metrics["membership_auc"],
+        "membership_privacy_score": mia_metrics["privacy_score"]
+    })
+    
+    # Attribute Inference Attack
+    if sensitive_attrs_train is not None:
+        X_train, _ = train_data
+        aia_metrics = evaluate_attribute_inference(model, X_train, sensitive_attrs_train, device)
+        results.update({
+            "attribute_auc": aia_metrics["attribute_auc"],
+            "attribute_accuracy": aia_metrics["attribute_accuracy"],
+            "attribute_privacy_score": aia_metrics["privacy_score"]
+        })
+    
+    # Overall privacy score
+    privacy_scores = [v for k, v in results.items() if 'privacy_score' in k]
+    results["overall_privacy_score"] = np.mean(privacy_scores) if privacy_scores else 0.0
+    
+    return results
+
+
+# ============================================================================
+# PERFORMANCE METRICS
+# ============================================================================
+
+def evaluate_performance(
+    model: nn.Module,
+    test_data: Tuple[np.ndarray, np.ndarray],
+    device: str = "cpu"
+) -> Dict[str, float]:
+    """
+    Evaluate model performance metrics.
+    
+    Returns:
+        Dictionary with accuracy, precision, recall, f1
+    """
+    model.eval()
+    X, y = test_data
+    X_tensor = torch.FloatTensor(X).to(device)
+    y_tensor = torch.LongTensor(y).to(device)
+    
+    with torch.no_grad():
+        output = model(X_tensor)
+        if isinstance(output, tuple):
+            logits = output[0]
+        else:
+            logits = output
+        
+        predictions = logits.argmax(dim=1).cpu().numpy()
+    
+    accuracy = accuracy_score(y, predictions)
+    precision, recall, f1, _ = precision_recall_fscore_support(
+        y, predictions, average='weighted', zero_division=0
+    )
+    
+    return {
+        "accuracy": accuracy,
+        "precision": precision,
+        "recall": recall,
+        "f1_score": f1
+    }