sensitivity.py

"""
Layer Sensitivity Analysis for Neural Networks

Sensitivity measures how much a layer's parameters affect model output.
High sensitivity → Important layer → Keep precision
Low sensitivity → Less critical → Can compress aggressively

Method: Perturbation analysis
- Add controlled noise to layer weights
- Measure change in model output
- Sensitivity = Output variance / Input variance

This helps identify which layers are most critical to model performance.
"""

import numpy as np
from typing import Dict, List, Tuple, Any, Optional
import torch
import torch.nn as nn


class SensitivityAnalyzer:
    """
    Perturbation-based sensitivity analysis for neural network layers.
    
    Measures how sensitive model outputs are to perturbations in each layer.
    This identifies critical vs non-critical layers for compression decisions.
    """
    
    def __init__(
        self,
        noise_level: float = 0.01,
        num_samples: int = 10,
        metric: str = 'mse'
    ):
        """
        Initialize sensitivity analyzer.
        
        Args:
            noise_level: Magnitude of perturbation (default: 1% of weight magnitude)
            num_samples: Number of perturbation samples to average
            metric: Sensitivity metric ('mse', 'cosine', 'kl')
        """
        self.noise_level = noise_level
        self.num_samples = num_samples
        self.metric = metric
        self.sensitivity_scores = {}
    
    def analyze_layer_sensitivity(
        self,
        model: nn.Module,
        layer_name: str,
        test_inputs: torch.Tensor,
        test_labels: Optional[torch.Tensor] = None
    ) -> float:
        """
        Measure sensitivity of a specific layer.
        
        Algorithm:
        1. Get baseline output: y_orig = model(x)
        2. For n samples:
           a. Perturb layer weights: w' = w + ε·N(0,1)
           b. Get perturbed output: y_pert = model(x)
           c. Measure difference: Δy = |y_orig - y_pert|
        3. Sensitivity = average(Δy)
        
        Args:
            model: PyTorch model
            layer_name: Name of layer to analyze
            test_inputs: Test input batch [batch, ...]
            test_labels: Optional labels for classification sensitivity
        
        Returns:
            sensitivity: Sensitivity score (0-1, higher = more sensitive)
        
        Example:
            >>> model = SimpleCNN()
            >>> inputs = torch.randn(32, 3, 32, 32)
            >>> sens_conv1 = analyzer.analyze_layer_sensitivity(model, 'conv1', inputs)
            >>> sens_conv2 = analyzer.analyze_layer_sensitivity(model, 'conv2', inputs)
            >>> print(f"conv1 sensitivity: {sens_conv1:.4f}")
            >>> print(f"conv2 sensitivity: {sens_conv2:.4f}")
            conv1 sensitivity: 0.8521  # High = critical layer
            conv2 sensitivity: 0.3142  # Lower = less critical
        """
        model.eval()
        
        # Find target layer
        target_module = None
        for name, module in model.named_modules():
            if name == layer_name:
                target_module = module
                break
        
        if target_module is None:
            raise ValueError(f"Layer '{layer_name}' not found in model")
        
        # Get baseline output
        with torch.no_grad():
            baseline_output = model(test_inputs)
        
        # Store original weights
        if not hasattr(target_module, 'weight'):
            return 0.0  # No weights to perturb
        
        original_weight = target_module.weight.data.clone()
        
        # Perturbation analysis
        sensitivities = []
        
        for _ in range(self.num_samples):
            # Add Gaussian noise
            noise = torch.randn_like(original_weight) * self.noise_level * original_weight.abs().mean()
            target_module.weight.data = original_weight + noise
            
            # Get perturbed output
            with torch.no_grad():
                perturbed_output = model(test_inputs)
            
            # Measure difference
            if self.metric == 'mse':
                diff = torch.mean((baseline_output - perturbed_output) ** 2).item()
            elif self.metric == 'cosine':
                # Cosine distance
                cos_sim = torch.nn.functional.cosine_similarity(
                    baseline_output.flatten(),
                    perturbed_output.flatten(),
                    dim=0
                ).item()
                diff = 1.0 - cos_sim
            elif self.metric == 'kl':
                # KL divergence (for classification)
                baseline_probs = torch.softmax(baseline_output, dim=-1)
                perturbed_probs = torch.softmax(perturbed_output, dim=-1)
                diff = torch.nn.functional.kl_div(
                    perturbed_probs.log(),
                    baseline_probs,
                    reduction='batchmean'
                ).item()
            else:
                diff = torch.mean(torch.abs(baseline_output - perturbed_output)).item()
            
            sensitivities.append(diff)
        
        # Restore original weights
        target_module.weight.data = original_weight
        
        # Average sensitivity
        sensitivity = np.mean(sensitivities)
        
        # Normalize to [0, 1] range (heuristic)
        sensitivity_normalized = min(sensitivity / (self.noise_level + 1e-6), 1.0)
        
        return sensitivity_normalized
    
    def profile_all_layers(
        self,
        model: nn.Module,
        test_inputs: torch.Tensor,
        layer_types: Optional[List[type]] = None
    ) -> Dict[str, float]:
        """
        Profile sensitivity of all layers in model.
        
        Args:
            model: PyTorch model
            test_inputs: Test input batch
            layer_types: Types of layers to profile (None = Conv2d and Linear)
        
        Returns:
            sensitivity_map: Dict mapping layer names to sensitivity scores
        
        Example:
            >>> model = SimpleCNN()
            >>> inputs = torch.randn(32, 3, 32, 32)
            >>> sens_map = analyzer.profile_all_layers(model, inputs)
            >>> for name, score in sorted(sens_map.items(), key=lambda x: -x[1]):
            >>>     print(f"{name}: {score:.4f}")
            conv1: 0.8521  # Most sensitive
            conv2: 0.6342
            fc:    0.2145  # Least sensitive
        """
        if layer_types is None:
            layer_types = [nn.Conv2d, nn.Linear]
        
        sensitivity_map = {}
        
        print(f"\n📊 Profiling layer sensitivity...")
        
        for name, module in model.named_modules():
            if any(isinstance(module, t) for t in layer_types):
                print(f"   Analyzing {name}...", end=' ')
                
                try:
                    sensitivity = self.analyze_layer_sensitivity(
                        model, name, test_inputs
                    )
                    sensitivity_map[name] = sensitivity
                    print(f"✓ {sensitivity:.4f}")
                except Exception as e:
                    print(f"✗ Error: {e}")
        
        # Store for later access
        self.sensitivity_scores = sensitivity_map
        
        return sensitivity_map
    
    def classify_sensitivity(self, sensitivity: float) -> str:
        """
        Classify layer by sensitivity level.
        
        Args:
            sensitivity: Sensitivity score (0-1)
        
        Returns:
            classification: 'critical', 'important', 'moderate', 'low'
        
        Thresholds:
            - Critical (>0.7): Essential layer, keep full precision
            - Important (0.4-0.7): Significant, use conservative compression
            - Moderate (0.2-0.4): Standard compression
            - Low (<0.2): Minimal impact, aggressive compression
        """
        if sensitivity > 0.7:
            return 'critical'
        elif sensitivity > 0.4:
            return 'important'
        elif sensitivity > 0.2:
            return 'moderate'
        else:
            return 'low'
    
    def get_compression_budget(self, sensitivity: float) -> int:
        """
        Get recommended bit-width based on sensitivity.
        
        Args:
            sensitivity: Sensitivity score (0-1)
        
        Returns:
            bits: Recommended quantization bit-width
        
        Mapping:
            - Critical (>0.7):   16 bits (FP16)
            - Important (0.4-0.7): 8 bits (INT8)
            - Moderate (0.2-0.4): 4 bits (INT4)
            - Low (<0.2):        2 bits (Binary/Ternary)
        """
        if sensitivity > 0.7:
            return 16  # FP16
        elif sensitivity > 0.4:
            return 8   # INT8
        elif sensitivity > 0.2:
            return 4   # INT4
        else:
            return 2   # Binary/Ternary
    
    def compute_gradient_sensitivity(
        self,
        model: nn.Module,
        loss_fn: nn.Module,
        test_inputs: torch.Tensor,
        test_labels: torch.Tensor
    ) -> Dict[str, float]:
        """
        Compute sensitivity using gradient information.
        
        This is faster than perturbation-based analysis and can be more
        accurate for differentiable models.
        
        Sensitivity = ||∂L/∂w||₂ (gradient norm)
        
        Args:
            model: PyTorch model
            loss_fn: Loss function
            test_inputs: Test inputs
            test_labels: Test labels
        
        Returns:
            gradient_sensitivity: Dict mapping layer names to gradient norms
        """
        model.train()  # Enable gradients
        
        # Forward pass
        outputs = model(test_inputs)
        loss = loss_fn(outputs, test_labels)
        
        # Backward pass
        model.zero_grad()
        loss.backward()
        
        # Collect gradient norms
        gradient_sensitivity = {}
        
        for name, module in model.named_modules():
            if isinstance(module, (nn.Conv2d, nn.Linear)):
                if hasattr(module, 'weight') and module.weight.grad is not None:
                    grad_norm = module.weight.grad.norm().item()
                    gradient_sensitivity[name] = grad_norm
        
        # Normalize
        max_grad = max(gradient_sensitivity.values()) if gradient_sensitivity else 1.0
        gradient_sensitivity = {
            k: v / max_grad for k, v in gradient_sensitivity.items()
        }
        
        model.eval()
        
        return gradient_sensitivity


def analyze_layer_sensitivity(
    model: nn.Module,
    layer_name: str,
    test_inputs: torch.Tensor,
    noise_level: float = 0.01,
    num_samples: int = 10
) -> float:
    """
    Convenience function to analyze single layer sensitivity.
    
    Args:
        model: PyTorch model
        layer_name: Name of layer to analyze
        test_inputs: Test input batch
        noise_level: Perturbation magnitude (default: 1%)
        num_samples: Number of perturbation samples
    
    Returns:
        sensitivity: Sensitivity score (0-1)
    """
    analyzer = SensitivityAnalyzer(
        noise_level=noise_level,
        num_samples=num_samples
    )
    return analyzer.analyze_layer_sensitivity(model, layer_name, test_inputs)


def create_sensitivity_report(
    sensitivity_map: Dict[str, float],
    save_path: Optional[str] = None
) -> str:
    """
    Create human-readable sensitivity analysis report.
    
    Args:
        sensitivity_map: Dict mapping layer names to sensitivity scores
        save_path: Path to save report (optional)
    
    Returns:
        report: Formatted report string
    """
    analyzer = SensitivityAnalyzer()
    
    report_lines = []
    report_lines.append("=" * 60)
    report_lines.append("LAYER SENSITIVITY ANALYSIS REPORT")
    report_lines.append("=" * 60)
    report_lines.append("")
    
    # Sort by sensitivity (descending)
    sorted_layers = sorted(
        sensitivity_map.items(),
        key=lambda x: -x[1]
    )
    
    report_lines.append(f"{'Layer':<20} {'Sensitivity':<12} {'Class':<12} {'Recommended'}")
    report_lines.append("-" * 60)
    
    for name, sensitivity in sorted_layers:
        classification = analyzer.classify_sensitivity(sensitivity)
        bits = analyzer.get_compression_budget(sensitivity)
        
        report_lines.append(
            f"{name:<20} {sensitivity:>8.4f}    {classification:<12} {bits}-bit"
        )
    
    report_lines.append("")
    report_lines.append("Classification Legend:")
    report_lines.append("  Critical  (>0.7):   Keep full precision (FP16)")
    report_lines.append("  Important (0.4-0.7): Conservative compression (INT8)")
    report_lines.append("  Moderate  (0.2-0.4): Standard compression (INT4)")
    report_lines.append("  Low       (<0.2):    Aggressive compression (Binary)")
    
    report = "\n".join(report_lines)
    
    if save_path:
        with open(save_path, 'w') as f:
            f.write(report)
        print(f"📄 Sensitivity report saved to: {save_path}")
    
    return report


if __name__ == '__main__':
    print("=" * 60)
    print("Layer Sensitivity Analysis")
    print("=" * 60)
    print("\nThis module implements perturbation-based sensitivity analysis:")
    print("  ✓ Measures layer importance via weight perturbation")
    print("  ✓ Identifies critical vs non-critical layers")
    print("  ✓ Guides precision allocation decisions")
    print("\nSensitivity Classification:")
    print("  Critical  (>0.7):   Essential → Keep full precision")
    print("  Important (0.4-0.7): Significant → Conservative compression")
    print("  Moderate  (0.2-0.4): Standard → INT4/INT8 compression")
    print("  Low       (<0.2):    Minimal impact → Aggressive compression")
    print("\nUsage:")
    print("  from passes.information_analysis import SensitivityAnalyzer")
    print("  analyzer = SensitivityAnalyzer()")
    print("  sensitivity_map = analyzer.profile_all_layers(model, test_inputs)")
    print("  report = create_sensitivity_report(sensitivity_map)")