mixed_precision_example.py

"""
ITC Mixed-Precision Quantization - Complete Example

This example demonstrates the full ITC pipeline with mixed-precision quantization:
1. Load a PyTorch model (ResNet18)
2. Analyze information density using IDM
3. Allocate precision adaptively per layer
4. Apply mixed-precision quantization
5. Validate accuracy preservation
6. Export optimized model

Author: ITC Team
Date: October 17, 2025
"""

import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import torch
import torch.nn as nn
import torchvision.models as models
from passes.information_analysis.entropy import EntropyAnalyzer
from passes.information_analysis.sensitivity import SensitivityAnalyzer
from passes.information_analysis.redundancy import RedundancyDetector
from passes.information_analysis.density import InformationDensityMetric
from passes.adaptive_precision import AdaptivePrecisionAllocator
from passes.mixed_precision import MixedPrecisionQuantizer, estimate_model_size
from export.onnx_exporter import export_to_onnx


def print_section(title):
    """Print formatted section header."""
    print("\n" + "=" * 80)
    print(f"  {title}")
    print("=" * 80)


def create_sample_model():
    """
    Create a simple CNN for demonstration.
    Using a smaller model than ResNet18 for faster execution.
    """
    class SimpleCNN(nn.Module):
        def __init__(self):
            super(SimpleCNN, self).__init__()
            self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
            self.bn1 = nn.BatchNorm2d(64)
            self.relu1 = nn.ReLU()
            
            self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
            self.bn2 = nn.BatchNorm2d(128)
            self.relu2 = nn.ReLU()
            
            self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
            self.bn3 = nn.BatchNorm2d(256)
            self.relu3 = nn.ReLU()
            
            self.pool = nn.AdaptiveAvgPool2d((1, 1))
            self.fc = nn.Linear(256, 10)
        
        def forward(self, x):
            x = self.relu1(self.bn1(self.conv1(x)))
            x = self.relu2(self.bn2(self.conv2(x)))
            x = self.relu3(self.bn3(self.conv3(x)))
            x = self.pool(x)
            x = torch.flatten(x, 1)
            x = self.fc(x)
            return x
    
    return SimpleCNN()


def main():
    """Main execution."""
    
    print("""
========================================================================
                                                                       
         ITC Mixed-Precision Quantization - Complete Example          
                                                                       
   Demonstrating information-theoretic compression with adaptive      
   precision allocation for optimal model compression.                
                                                                       
========================================================================
    """)
    
    # Configuration
    input_shape = (1, 3, 32, 32)  # CIFAR-10 size
    compression_target = 6.0  # Target 6x compression
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    print(f"Configuration:")
    print(f"  • Input shape: {input_shape}")
    print(f"  • Compression target: {compression_target}x")
    print(f"  • Device: {device}")
    
    # Step 1: Create/Load Model
    print_section("Step 1: Creating Model")
    
    model = create_sample_model()
    model.eval()
    
    # Count parameters
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    print(f"✓ Model created: SimpleCNN")
    print(f"  • Total parameters: {total_params:,}")
    print(f"  • Trainable parameters: {trainable_params:,}")
    print(f"  • Layers: {len(list(model.named_modules())):,}")
    
    # Step 2: Analyze Information Density
    print_section("Step 2: Analyzing Information Density (IDM)")
    
    print("The Information Density Metric (IDM) combines three key signals:")
    print("  1. Entropy: Information content in layer weights")
    print("  2. Sensitivity: Layer criticality via perturbation analysis")
    print("  3. Redundancy: Detection of redundant filters/channels")
    print("\nIDM Formula: (Entropy × Sensitivity) / (Size × (1 + Redundancy))\n")
    
    # Step 2a: Compute entropy
    print("[2a] Computing entropy for each layer...")
    entropy_analyzer = EntropyAnalyzer()
    entropy_map = entropy_analyzer.analyze_weight_entropy(model)
    print(f"✓ Entropy analysis complete for {len(entropy_map)} layers")
    
    # Step 2b: Compute sensitivity
    print("\n[2b] Computing sensitivity (this may take a moment)...")
    sensitivity_analyzer = SensitivityAnalyzer()
    # Create test inputs for sensitivity analysis
    test_inputs = torch.randn(5, *input_shape[1:])  # 5 samples
    sensitivity_map = sensitivity_analyzer.profile_all_layers(model, test_inputs)
    print(f"✓ Sensitivity analysis complete for {len(sensitivity_map)} layers")
    
    # Step 2c: Detect redundancy
    print("\n[2c] Detecting redundancy...")
    redundancy_detector = RedundancyDetector()
    redundancy_map = redundancy_detector.compute_redundancy_score(model)
    print(f"✓ Redundancy analysis complete for {len(redundancy_map)} layers")
    
    # Step 2d: Compute IDM
    print("\n[2d] Computing Information Density Metric...")
    idm_analyzer = InformationDensityMetric()
    idm_results = idm_analyzer.analyze_model(model, entropy_map, sensitivity_map, redundancy_map)
    
    print(f"\n✓ IDM analysis complete for {len(idm_results)} layers\n")
    print(f"{'Layer':<20} {'IDM Score':<12} {'Classification':<15} {'Recommendation':<15}")
    print("-" * 62)
    
    for name, metrics in idm_results.items():
        idm = metrics['idm']
        classification = metrics['classification']
        recommendation = metrics.get('precision_recommendation', 'N/A')
        print(f"{name:<20} {idm:<12.6f} {classification:<15} {recommendation:<15}")
    
    # Interpretation
    print("\nInterpretation:")
    print("  • critical (IDM > 0.001):    Keep high precision (FP16)")
    print("  • important (IDM > 0.0001):  INT8 quantization")
    print("  • standard (IDM > 0.00001):  INT4 quantization")
    print("  • compressible (IDM < 0.00001): Binary quantization")
    
    # Step 3: Allocate Precision
    print_section(f"Step 3: Allocating Precision (Target: {compression_target}x)")
    
    print("Allocating bit-widths per layer based on IDM scores...")
    
    allocator = AdaptivePrecisionAllocator(compression_target=compression_target)
    precision_map = allocator.allocate(model, idm_results)
    
    print(f"\n✓ Precision allocation complete\n")
    print(f"{'Layer':<20} {'Precision':<12} {'Format':<15}")
    print("-" * 47)
    
    for name, spec in precision_map.items():
        bits = spec['bits']
        dtype = spec['dtype']
        print(f"{name:<20} {bits}-bit       {dtype:<15}")
    
    # Count precision distribution
    precision_counts = {}
    for name, spec in precision_map.items():
        bits = spec['bits']
        precision_counts[bits] = precision_counts.get(bits, 0) + 1
    
    print(f"\nPrecision Distribution:")
    for bits in sorted(precision_counts.keys(), reverse=True):
        count = precision_counts[bits]
        percentage = (count / len(precision_map)) * 100
        print(f"  • {bits}-bit: {count} layers ({percentage:.1f}%)")
    
    # Step 4: Apply Quantization
    print_section("Step 4: Applying Mixed-Precision Quantization")
    
    print("Quantizing model with allocated precision schedule...")
    
    quantizer = MixedPrecisionQuantizer()
    quantized_model = quantizer.quantize_model(model, precision_map)
    
    print(f"✓ Quantization applied to {len(precision_map)} layers")
    
    # Estimate model size
    size_info = estimate_model_size(model, precision_map)
    
    print(f"\nModel Size Analysis:")
    print(f"  • Original size: {size_info['size_fp32_mb']:.2f} MB")
    print(f"  • Quantized size: {size_info['size_quantized_mb']:.2f} MB")
    print(f"  • Compression ratio: {size_info['compression_ratio']:.2f}x")
    print(f"  • Space saved: {size_info['size_saved_percent']:.1f}%")
    
    # Step 5: Validate Accuracy
    print_section("Step 5: Validating Accuracy Preservation")
    
    print("Running inference on sample inputs...")
    
    # Create sample test data
    num_samples = 10
    test_inputs = torch.randn(num_samples, *input_shape[1:])
    
    # Original model inference
    model.eval()
    with torch.no_grad():
        original_outputs = model(test_inputs)
    
    # Quantized model inference
    quantized_model.eval()
    with torch.no_grad():
        quantized_outputs = quantized_model(test_inputs)
    
    # Calculate errors
    abs_diff = torch.abs(original_outputs - quantized_outputs)
    max_diff = torch.max(abs_diff).item()
    mean_diff = torch.mean(abs_diff).item()
    
    # Relative error
    original_mean = torch.mean(torch.abs(original_outputs)).item()
    relative_error = (mean_diff / original_mean) * 100 if original_mean > 0 else 0
    
    print(f"\n✓ Validation complete on {num_samples} samples\n")
    print(f"Error Metrics:")
    print(f"  • Maximum difference: {max_diff:.6f}")
    print(f"  • Mean difference: {mean_diff:.6f}")
    print(f"  • Relative error: {relative_error:.2f}%")
    
    # Quality assessment
    if relative_error < 1.0:
        quality = "EXCELLENT"
        emoji = "✓✓✓"
    elif relative_error < 3.0:
        quality = "VERY GOOD"
        emoji = "✓✓"
    elif relative_error < 5.0:
        quality = "GOOD"
        emoji = "✓"
    elif relative_error < 10.0:
        quality = "ACCEPTABLE"
        emoji = "⚠"
    else:
        quality = "POOR"
        emoji = "✗"
    
    print(f"\nQuality Assessment: {emoji} {quality}")
    
    # Step 6: Export Model
    print_section("Step 6: Exporting Optimized Model")
    
    output_dir = "./outputs/mixed_precision_example"
    os.makedirs(output_dir, exist_ok=True)
    
    # Export to ONNX
    onnx_path = os.path.join(output_dir, "quantized_model.onnx")
    print(f"Exporting to ONNX format: {onnx_path}")
    
    try:
        export_to_onnx(quantized_model, input_shape, onnx_path)
        print(f"✓ ONNX export successful")
    except Exception as e:
        print(f"✗ ONNX export failed: {e}")
    
    # Save precision schedule
    schedule_path = os.path.join(output_dir, "precision_schedule.txt")
    print(f"\nSaving precision schedule: {schedule_path}")
    
    with open(schedule_path, 'w') as f:
        f.write("ITC Mixed-Precision Quantization - Precision Schedule\n")
        f.write("=" * 60 + "\n\n")
        for name, spec in precision_map.items():
            f.write(f"{name}:\n")
            f.write(f"  Bits: {spec['bits']}\n")
            f.write(f"  Dtype: {spec['dtype']}\n")
            f.write(f"  Method: {spec['method']}\n")
            f.write(f"  Description: {spec['description']}\n\n")
    
    print(f"✓ Precision schedule saved")
    
    # Final Summary
    print_section("Summary")
    
    print(f"""
✓ ITC Mixed-Precision Quantization Complete!

Key Results:
  • Compression achieved: {size_info['compression_ratio']:.2f}x (target: {compression_target}x)
  • Model size: {size_info['size_fp32_mb']:.2f} MB → {size_info['size_quantized_mb']:.2f} MB
  • Space saved: {size_info['size_saved_percent']:.1f}%
  • Accuracy loss: {relative_error:.2f}% ({quality})
  • Layers quantized: {len(precision_map)}
  
Outputs:
  • Quantized ONNX model: {onnx_path}
  • Precision schedule: {schedule_path}

Next Steps:
  1. Deploy the quantized ONNX model for inference
  2. Integrate with ONNX Runtime or TensorRT for optimal performance
  3. Fine-tune the model if accuracy needs improvement
  4. Adjust compression target for different size/accuracy trade-offs
    """)
    
    print("=" * 80)
    print("\n🎉 Example completed successfully!\n")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n⚠ Execution interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n✗ Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)