# Install PyTorch with CUDA support
pip install torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118

# Install optimization libraries
pip install torch-pruning==1.3.3  # Model pruning
pip install onnx==1.15.0 onnxruntime==1.16.0  # ONNX conversion
pip install tensorrt==8.6.1  # NVIDIA TensorRT (requires CUDA)

# Install utilities
pip install thop==0.1.1  # FLOPs calculation
pip install matplotlib==3.8.0 seaborn==0.12.2  # Visualization

# post_training_quantization.py - Quantize trained model without retraining

import torch
import torchvision.models as models
from torch.quantization import quantize_dynamic, quantize_static, get_default_qconfig
import torch.quantization

def dynamic_quantization_example():
    """
    Dynamic quantization: Quantize weights (static), activations (dynamic at runtime)
    Best for: RNNs, LSTMs, Transformers
    """
    # Load pre-trained model
    model = models.resnet18(pretrained=True)
    model.eval()

    # Apply dynamic quantization to Linear and LSTM layers
    quantized_model = quantize_dynamic(
        model,
        {torch.nn.Linear},  # Which layers to quantize
        dtype=torch.qint8
    )

    # Compare model sizes
    def get_model_size(model):
        torch.save(model.state_dict(), "temp.p")
        size = os.path.getsize("temp.p") / 1e6  # MB
        os.remove("temp.p")
        return size

    print(f"Original model: {get_model_size(model):.2f} MB")
    print(f"Quantized model: {get_model_size(quantized_model):.2f} MB")

    return quantized_model


def static_quantization_example(model, calibration_dataloader):
    """
    Static quantization: Quantize both weights and activations (static)
    Requires calibration data to compute activation ranges
    Best for: CNNs on CPU
    """
    # Set model to evaluation mode
    model.eval()

    # Fuse Conv+BatchNorm+ReLU layers for better performance
    model_fused = torch.quantization.fuse_modules(
        model,
        [['conv1', 'bn1', 'relu']],  # Specify layers to fuse
        inplace=False
    )

    # Specify quantization config
    model_fused.qconfig = get_default_qconfig('fbgemm')  # x86 CPU backend

    # Prepare model for static quantization
    model_prepared = torch.quantization.prepare(model_fused, inplace=False)

    # Calibrate with representative data
    print("Calibrating quantization parameters...")
    with torch.no_grad():
        for inputs, _ in calibration_dataloader:
            model_prepared(inputs)

    # Convert to quantized model
    quantized_model = torch.quantization.convert(model_prepared, inplace=False)

    return quantized_model


# Benchmark inference speed
def benchmark_model(model, input_tensor, num_iterations=100):
    """
    Measure inference latency
    """
    import time

    model.eval()
    with torch.no_grad():
        # Warmup
        for _ in range(10):
            _ = model(input_tensor)

        # Measure
        start = time.time()
        for _ in range(num_iterations):
            _ = model(input_tensor)
        end = time.time()

    avg_time = (end - start) / num_iterations * 1000  # ms
    return avg_time


# Example usage
if __name__ == "__main__":
    # Create dummy data
    input_tensor = torch.randn(1, 3, 224, 224)

    # Original model
    model_fp32 = models.resnet18(pretrained=True)
    model_fp32.eval()

    # Dynamic quantization
    model_int8 = dynamic_quantization_example()

    # Benchmark
    latency_fp32 = benchmark_model(model_fp32, input_tensor)
    latency_int8 = benchmark_model(model_int8, input_tensor)

    print(f"\nInference Latency (CPU):")
    print(f"FP32: {latency_fp32:.2f} ms")
    print(f"INT8: {latency_int8:.2f} ms")
    print(f"Speedup: {latency_fp32 / latency_int8:.2f}x")

# quantization_aware_training.py - Train with quantization in mind

import torch
import torch.nn as nn
import torch.quantization

class SimpleConvNet(nn.Module):
    """
    Example CNN for QAT demonstration
    """
    def __init__(self, num_classes=10):
        super().__init__()
        self.quant = torch.quantization.QuantStub()  # Quantize input
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU()
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(64, num_classes)
        self.dequant = torch.quantization.DeQuantStub()  # Dequantize output

    def forward(self, x):
        x = self.quant(x)
        x = self.relu1(self.bn1(self.conv1(x)))
        x = self.relu2(self.bn2(self.conv2(x)))
        x = self.pool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        x = self.dequant(x)
        return x


def train_with_quantization_awareness(model, train_loader, val_loader, epochs=10):
    """
    QAT training loop
    """
    # Fuse layers
    model.train()
    model.fuse_model()  # Implement custom fuse_model() method

    # Set QAT configuration
    model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')

    # Prepare for QAT
    model_prepared = torch.quantization.prepare_qat(model, inplace=False)

    # Training loop
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model_prepared.parameters(), lr=0.001)

    for epoch in range(epochs):
        model_prepared.train()
        for inputs, labels in train_loader:
            optimizer.zero_grad()
            outputs = model_prepared(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

        # Validate
        model_prepared.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                outputs = model_prepared(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = 100 * correct / total
        print(f'Epoch {epoch+1}/{epochs}, Accuracy: {accuracy:.2f}%')

    # Convert to quantized model
    model_prepared.eval()
    quantized_model = torch.quantization.convert(model_prepared, inplace=False)

    return quantized_model

# knowledge_distillation.py - Distill knowledge from teacher to student

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    """
    Combined loss for knowledge distillation
    Loss = alpha * KL(teacher, student) + (1-alpha) * CE(student, labels)
    """
    def __init__(self, temperature=3.0, alpha=0.7):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha
        self.ce_loss = nn.CrossEntropyLoss()

    def forward(self, student_logits, teacher_logits, labels):
        # Soft targets from teacher (with temperature)
        soft_targets = F.softmax(teacher_logits / self.temperature, dim=1)
        soft_prob = F.log_softmax(student_logits / self.temperature, dim=1)

        # Distillation loss (KL divergence)
        distillation_loss = F.kl_div(
            soft_prob,
            soft_targets,
            reduction='batchmean'
        ) * (self.temperature ** 2)

        # Student loss (cross-entropy with true labels)
        student_loss = self.ce_loss(student_logits, labels)

        # Combined loss
        return self.alpha * distillation_loss + (1 - self.alpha) * student_loss


def distill_model(teacher_model, student_model, train_loader, val_loader, epochs=20):
    """
    Train student model to mimic teacher model
    """
    teacher_model.eval()  # Teacher in eval mode
    student_model.train()

    criterion = DistillationLoss(temperature=3.0, alpha=0.7)
    optimizer = torch.optim.Adam(student_model.parameters(), lr=0.001)

    for epoch in range(epochs):
        student_model.train()
        total_loss = 0

        for inputs, labels in train_loader:
            optimizer.zero_grad()

            # Get teacher predictions (no gradient)
            with torch.no_grad():
                teacher_logits = teacher_model(inputs)

            # Get student predictions
            student_logits = student_model(inputs)

            # Calculate distillation loss
            loss = criterion(student_logits, teacher_logits, labels)

            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        # Validation
        student_model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                outputs = student_model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = 100 * correct / total
        avg_loss = total_loss / len(train_loader)
        print(f'Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}, Accuracy: {accuracy:.2f}%')

    return student_model


# Example: Distill ResNet-50 to ResNet-18
if __name__ == "__main__":
    import torchvision.models as models

    # Teacher: Large model (ResNet-50)
    teacher = models.resnet50(pretrained=True)
    teacher.eval()

    # Student: Small model (ResNet-18)
    student = models.resnet18(pretrained=False)

    # Assuming you have train_loader and val_loader
    # student_distilled = distill_model(teacher, student, train_loader, val_loader)

    # Compare model sizes
    teacher_params = sum(p.numel() for p in teacher.parameters())
    student_params = sum(p.numel() for p in student.parameters())

    print(f"Teacher parameters: {teacher_params:,} ({teacher_params/1e6:.1f}M)")
    print(f"Student parameters: {student_params:,} ({student_params/1e6:.1f}M)")
    print(f"Compression ratio: {teacher_params/student_params:.2f}x")

# pruning.py - Structured and unstructured pruning

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

def unstructured_pruning_example(model, pruning_amount=0.3):
    """
    Magnitude-based unstructured pruning
    Remove individual weights with smallest magnitudes
    """
    # Prune 30% of connections in each Conv2d layer
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            prune.l1_unstructured(module, name='weight', amount=pruning_amount)

            # Make pruning permanent
            prune.remove(module, 'weight')

    return model


def structured_pruning_example(model, pruning_amount=0.2):
    """
    Structured pruning: Remove entire channels/filters
    Better for actual speedup (unlike unstructured which needs sparse operations)
    """
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            # Prune 20% of output channels (entire filters)
            prune.ln_structured(
                module,
                name='weight',
                amount=pruning_amount,
                n=2,  # L2 norm
                dim=0  # Prune along output channel dimension
            )

            prune.remove(module, 'weight')

    return model


def iterative_pruning(model, train_loader, val_loader, initial_sparsity=0.0, final_sparsity=0.7, epochs_per_step=5, pruning_steps=10):
    """
    Iterative magnitude pruning (IMP)
    Gradually increase sparsity while retraining
    """
    import numpy as np

    # Calculate sparsity schedule
    sparsity_levels = np.linspace(initial_sparsity, final_sparsity, pruning_steps)

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    for step, target_sparsity in enumerate(sparsity_levels):
        print(f"\nPruning step {step+1}/{pruning_steps}, Target sparsity: {target_sparsity:.2%}")

        # Apply pruning
        for module in model.modules():
            if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
                prune.l1_unstructured(module, name='weight', amount=target_sparsity)

        # Fine-tune model
        for epoch in range(epochs_per_step):
            model.train()
            for inputs, labels in train_loader:
                optimizer.zero_grad()
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()

        # Evaluate
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = 100 * correct / total
        print(f"Accuracy after pruning: {accuracy:.2f}%")

    # Make pruning permanent
    for module in model.modules():
        if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
            if hasattr(module, 'weight_mask'):
                prune.remove(module, 'weight')

    return model


def calculate_sparsity(model):
    """
    Calculate percentage of zero weights
    """
    total_params = 0
    zero_params = 0

    for param in model.parameters():
        total_params += param.numel()
        zero_params += (param == 0).sum().item()

    sparsity = 100 * zero_params / total_params
    return sparsity

# onnx_conversion.py - Convert PyTorch to ONNX format

import torch
import torch.onnx
import onnx
import onnxruntime

def convert_to_onnx(model, input_shape, onnx_path="model.onnx"):
    """
    Convert PyTorch model to ONNX format
    """
    model.eval()

    # Create dummy input
    dummy_input = torch.randn(input_shape)

    # Export to ONNX
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        export_params=True,
        opset_version=14,
        do_constant_folding=True,  # Constant folding optimization
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={
            'input': {0: 'batch_size'},
            'output': {0: 'batch_size'}
        }
    )

    # Verify ONNX model
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)

    print(f"Model converted to ONNX: {onnx_path}")

    return onnx_path


def benchmark_onnx(onnx_path, input_shape, num_iterations=100):
    """
    Benchmark ONNX Runtime inference
    """
    import time
    import numpy as np

    # Load ONNX model
    session = onnxruntime.InferenceSession(onnx_path)

    # Create dummy input
    input_data = np.random.randn(*input_shape).astype(np.float32)

    # Warmup
    for _ in range(10):
        _ = session.run(None, {'input': input_data})

    # Benchmark
    start = time.time()
    for _ in range(num_iterations):
        _ = session.run(None, {'input': input_data})
    end = time.time()

    avg_time = (end - start) / num_iterations * 1000  # ms
    return avg_time

# tensorrt_optimization.sh - Convert ONNX to TensorRT engine

# Install TensorRT (Ubuntu)
# Download from https://developer.nvidia.com/tensorrt

# Convert ONNX to TensorRT with FP16 precision
trtexec --onnx=model.onnx \
        --saveEngine=model_fp16.trt \
        --fp16 \
        --workspace=4096  # 4GB workspace

# Convert with INT8 quantization (requires calibration data)
trtexec --onnx=model.onnx \
        --saveEngine=model_int8.trt \
        --int8 \
        --calib=calibration_cache.bin

# Benchmark TensorRT engine
trtexec --loadEngine=model_fp16.trt --iterations=1000

# Expected output:
# Latency: min = 2.5 ms, max = 3.1 ms, mean = 2.7 ms, median = 2.6 ms