Upload mia.py

mia.py

@@ -0,0 +1,203 @@
+"""
+MIA Calibrator for the Dependable FungAI Project.
+
+This script performs a post-training analysis to determine the optimal
+threshold for Membership Inference Attacks (MIA). It collects the two most
+reliable metrics: Confidence Score and Loss Value.
+
+It uses a balanced logistic regression model on confidence scores to
+determine a single, robust membership threshold.
+
+The script outputs:
+1.  A visualization of the confidence score distributions.
+2.  An updated TBOM.json file with a detailed 'membership_inference_analysis' key,
+    containing the calculated threshold and summary statistics for the collected metrics.
+
+This script should be run AFTER TBOM.py has successfully completed.
+"""
+
+import os
+import json
+import torch
+import torch.nn as nn
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import roc_auc_score
+from torchvision import datasets
+from torch.utils.data import DataLoader, Subset
+import warnings
+
+# Suppress warnings for clean output
+warnings.filterwarnings('ignore')
+
+# --- Import from existing project files ---
+from TBOM import HybridMLP
+from IBOM import DetailedIBOMGenerator, get_paths_with_smart_fallback
+
+def get_mia_metrics(model_wrapper, dataset, batch_size=32):
+    """
+    Runs inference on a dataset and returns a list of dictionaries,
+    each containing MIA metrics (confidence, loss) for a sample.
+    """
+    print(f"Calculating MIA metrics (Loss, Confidence) for {len(dataset)} samples...")
+    all_metrics = []
+    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+    criterion = nn.BCEWithLogitsLoss(reduction='none') # Use 'none' to get per-sample loss
+
+    model = model_wrapper.model
+    model.eval()
+
+    with torch.no_grad():
+        for i, (images, labels) in enumerate(data_loader):
+            images = images.to(model_wrapper.device)
+            labels_float = labels.unsqueeze(1).float().to(model_wrapper.device)
+
+            # --- Feature Extraction (mirrors IBOM) ---
+            img_features = model_wrapper.clip_model.encode_image(images).float()
+            img_features_norm = img_features / img_features.norm(dim=1, keepdim=True)
+            concept_scores = img_features_norm @ model_wrapper.text_embeddings.T
+            hybrid_features = torch.cat([img_features, concept_scores], dim=1)
+
+            # --- Metric Calculation ---
+            logits = model(hybrid_features)
+            probabilities = torch.sigmoid(logits)
+            
+            # 1. Confidence Score
+            confidences = torch.max(probabilities, 1 - probabilities).squeeze()
+            
+            # 2. Loss Value (per-sample)
+            losses = criterion(logits, labels_float).squeeze()
+
+            # --- Store metrics for each sample in the batch ---
+            for j in range(images.size(0)):
+                all_metrics.append({
+                    'confidence': confidences[j].item(),
+                    'loss': losses[j].item()
+                })
+
+            if (i + 1) % 50 == 0:
+                print(f"  Processed { (i + 1) * batch_size } / {len(dataset)} samples...")
+
+    return all_metrics
+
+def find_best_threshold_lr_balanced(member_scores, non_member_scores):
+    """
+    Trains a balanced logistic regression model on confidence scores to find an MIA threshold.
+    """
+    print("\nFinding optimal threshold using a BALANCED logistic regression attack model...")
+    X = np.concatenate([member_scores, non_member_scores]).reshape(-1, 1)
+    y = np.concatenate([np.ones_like(member_scores), np.zeros_like(non_member_scores)])
+    
+    attack_model = LogisticRegression(solver='liblinear', class_weight='balanced')
+    attack_model.fit(X, y)
+    
+    intercept = attack_model.intercept_[0]
+    coef = attack_model.coef_[0][0]
+    threshold = -intercept / coef if coef != 0 else 0.5
+        
+    attack_probs = attack_model.predict_proba(X)[:, 1]
+    attack_auc = roc_auc_score(y, attack_probs)
+    
+    print(f"BALANCED Attack Model Trained. Optimal Threshold: {threshold:.4f}, Attack AUC: {attack_auc:.4f}")
+    return float(threshold), float(attack_auc)
+
+def visualize_distributions(member_scores, non_member_scores, threshold, output_path):
+    """
+    Creates and saves a histogram plot of the confidence distributions.
+    """
+    print(f"\nGenerating and saving visualization to {output_path}...")
+    plt.style.use('seaborn-v0_8-whitegrid')
+    plt.figure(figsize=(12, 7))
+    
+    plt.hist(non_member_scores, bins=50, density=True, alpha=0.7, label='Non-Members (Test Set)', color='darkorange')
+    plt.hist(member_scores, bins=50, density=True, alpha=0.7, label='Members (Train Set)', color='royalblue')
+    
+    plt.axvline(threshold, color='crimson', linestyle='--', linewidth=2.5, label=f'Decision Threshold ({threshold:.3f})')
+    
+    plt.title('Confidence Score Distributions: Members vs. Non-Members', fontsize=16, fontweight='bold')
+    plt.xlabel('Model Prediction Confidence', fontsize=12)
+    plt.ylabel('Density', fontsize=12)
+    plt.legend(fontsize=11)
+    plt.xlim(0.5, 1.0)
+    plt.tight_layout()
+    
+    plt.savefig(output_path, dpi=300)
+    plt.close()
+
+def main():
+    """
+    Main function to orchestrate the MIA calibration process.
+    """
+    print("--- Starting Membership Inference Attack (MIA) Calibration ---")
+    
+    try:
+        tbom_path, model_path, csv_path = get_paths_with_smart_fallback()
+        output_dir = os.path.dirname(tbom_path)
+        
+        print(f"Loading data splits from {tbom_path}")
+        with open(tbom_path, 'r') as f:
+            tbom_data = json.load(f)
+        
+        train_val_indices = tbom_data['data_summary']['data_splits']['train_validation_set']['indices']
+        test_indices = tbom_data['data_summary']['data_splits']['test_set']['indices']
+        image_root = tbom_data['data_summary']['image_dataset_path']
+
+        print("\nInitializing model and data pipeline...")
+        ibom_generator = DetailedIBOMGenerator(model_path, tbom_path, csv_path)
+        
+        full_dataset = datasets.ImageFolder(root=image_root, transform=ibom_generator.preprocess)
+        member_dataset = Subset(full_dataset, train_val_indices)
+        non_member_dataset = Subset(full_dataset, test_indices)
+
+        member_metrics = get_mia_metrics(ibom_generator, member_dataset)
+        non_member_metrics = get_mia_metrics(ibom_generator, non_member_dataset)
+
+        member_confidences = np.array([m['confidence'] for m in member_metrics])
+        non_member_confidences = np.array([m['confidence'] for m in non_member_metrics])
+        
+        final_threshold, attack_auc = find_best_threshold_lr_balanced(member_confidences, non_member_confidences)
+        
+        print(f"\n--- Final Threshold Selected ---")
+        print(f"  Method: Balanced Logistic Regression")
+        print(f"  Threshold (on Confidence Score): {final_threshold:.4f}")
+
+        viz_path = os.path.join(output_dir, 'mia_confidence_distribution.png')
+        visualize_distributions(member_confidences, non_member_confidences, final_threshold, viz_path)
+        
+        print(f"\nUpdating {tbom_path} with detailed MIA results...")
+
+        member_df = pd.DataFrame(member_metrics)
+        non_member_df = pd.DataFrame(non_member_metrics)
+
+        tbom_data['membership_inference_analysis'] = {
+            'description': "Analysis to distinguish members from non-members using Loss and Confidence metrics.",
+            'threshold_finding_summary': {
+                'metric_used_for_threshold': 'confidence_score',
+                'method': "A balanced logistic regression was used to find the optimal threshold.",
+                'decision_threshold': final_threshold,
+                'attack_auc_score_on_confidence': attack_auc,
+            },
+            'metric_statistics': {
+                'members': member_df.describe().to_dict(),
+                'non_members': non_member_df.describe().to_dict()
+            },
+            'interpretation': "Clear separation in metric distributions (lower loss and higher confidence for members) indicates data memorization.",
+            'visualization_artifact': viz_path
+        }
+        
+        with open(tbom_path, 'w') as f:
+            json.dump(tbom_data, f, indent=4)
+            
+        print("\n--- MIA Calibration Complete! ---")
+        print(f"✅ TBOM file successfully updated with focused membership analysis.")
+        print(f"✅ Distribution plot saved to {viz_path}")
+
+    except Exception as e:
+        print(f"\n❌ An error occurred during MIA calibration: {e}")
+        import traceback
+        traceback.print_exc()
+
+if __name__ == "__main__":
+    main()