05 - Pruning Reti Neurali: Ridurre Parametri

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune
import numpy as np

# ===================================================================
# MAGNITUDE PRUNING CON PYTORCH NATIVE API
# ===================================================================

class ConvNet(nn.Module):
    """Modello CNN semplice per dimostrare il pruning."""
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(4)
        )
        self.classifier = nn.Sequential(
            nn.Linear(128 * 4 * 4, 256),
            nn.ReLU(),
            nn.Linear(256, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)


model = ConvNet()

# --- Pruning L1 non strutturato (magnitude-based) ---
# Rimuove il 30% dei pesi con valore assoluto minore
prune.l1_unstructured(
    model.features[0],   # Layer da pruning
    name='weight',       # Parametro da prunare
    amount=0.30          # Percentuale da rimuovere (30%)
)

# --- Pruning Random (baseline di confronto) ---
prune.random_unstructured(
    model.features[2],
    name='weight',
    amount=0.30
)

# --- Analisi sparsita risultante ---
def compute_sparsity(module):
    """Calcola la sparsita effettiva di un modulo."""
    total = 0
    zeros = 0
    for param in module.parameters():
        total += param.numel()
        zeros += (param == 0).sum().item()
    return zeros / total if total > 0 else 0.0

print("Sparsita Conv1:", f"{compute_sparsity(model.features[0]):.1%}")
print("Sparsita Conv2:", f"{compute_sparsity(model.features[2]):.1%}")

# --- Verifica la struttura interna del pruning ---
# PyTorch crea weight_orig (originale) + weight_mask (0/1)
print("\nParametri di model.features[0] dopo pruning:")
for name, param in model.features[0].named_parameters():
    print(f"  {name}: shape={param.shape}")
for name, buf in model.features[0].named_buffers():
    print(f"  buffer {name}: shape={buf.shape}")

# --- Rimozione della maschera (make permanent) ---
# Dopo retraining, si consolida: il modello torna a usare 'weight'
prune.remove(model.features[0], 'weight')
print("\nDopo prune.remove: parametri di model.features[0]:")
for name, _ in model.features[0].named_parameters():
    print(f"  {name}")

# --- Global Pruning: pruning globale su tutto il modello ---
# Più efficace del pruning per-layer: usa una soglia globale
parameters_to_prune = (
    (model.features[0], 'weight'),
    (model.features[2], 'weight'),
    (model.classifier[0], 'weight'),
    (model.classifier[2], 'weight'),
)

prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.40,  # Rimuove 40% globalmente (non per layer)
)

# Sparsita finale per layer
for module_name, module in model.named_modules():
    if isinstance(module, (nn.Conv2d, nn.Linear)):
        if hasattr(module, 'weight_mask'):
            sparsity = (module.weight_mask == 0).float().mean().item()
            print(f"{module_name}: sparsita {sparsity:.1%}")

# pip install torch-pruning

import torch
import torch.nn as nn
import torch_pruning as tp

# ===================================================================
# PRUNING STRUTTURATO CON TORCH-PRUNING
# ===================================================================

class ResidualBlock(nn.Module):
    """Blocco residuale: Torch-Pruning gestisce la skip connection automaticamente."""
    def __init__(self, channels=64):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(channels)

    def forward(self, x):
        residual = x
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        return self.relu(out + residual)  # Skip connection


class SimpleResNet(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.layer1 = ResidualBlock(64)
        self.layer2 = ResidualBlock(64)
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(64, num_classes)

    def forward(self, x):
        x = self.stem(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.pool(x).view(x.size(0), -1)
        return self.fc(x)


model = SimpleResNet()
model.eval()

# Input di esempio per tracciare le dipendenze
example_input = torch.randn(1, 3, 32, 32)

# --- Costruzione del grafo delle dipendenze ---
DG = tp.DependencyGraph().build_dependency(model, example_inputs=example_input)

# --- Analisi del modello PRIMA del pruning ---
macs_before, params_before = tp.utils.count_ops_and_params(model, example_input)
print(f"Parametri PRIMA: {params_before / 1e6:.2f}M")
print(f"MACs PRIMA: {macs_before / 1e9:.3f}G")

# --- Definizione della strategia di pruning ---
# Pruning per magnitudine L1 dei filtri (L2 disponibile con tp.strategy.L2Strategy)
pruner = tp.pruner.MagnitudePruner(
    model,
    example_inputs=example_input,
    importance=tp.importance.MagnitudeImportance(p=1),  # L1 norm
    iterative_steps=5,          # Pruning iterativo in 5 step
    ch_sparsity=0.5,            # Rimuove il 50% dei canali
    ignored_layers=[model.fc], # Non pruning il classificatore finale
)

# --- Esecuzione del pruning (un singolo step) ---
pruner.step()

# --- Analisi del modello DOPO il pruning ---
macs_after, params_after = tp.utils.count_ops_and_params(model, example_input)
print(f"\nParametri DOPO: {params_after / 1e6:.2f}M")
print(f"MACs DOPO: {macs_after / 1e9:.3f}G")
print(f"Riduzione parametri: {(1 - params_after/params_before):.1%}")
print(f"Riduzione MACs: {(1 - macs_after/macs_before):.1%}")

# --- Verifica architettura post-pruning ---
print("\nArchitettura post-pruning:")
for name, module in model.named_modules():
    if isinstance(module, nn.Conv2d):
        print(f"  {name}: Conv2d({module.in_channels}, {module.out_channels}, ...)")

# Output tipico:
# Parametri PRIMA: 0.15M | MACs PRIMA: 0.009G
# Parametri DOPO:  0.04M | MACs DOPO:  0.003G
# Riduzione parametri: 75% | Riduzione MACs: 72%
# layer1.conv1: Conv2d(32, 32, ...) <- da 64 a 32 canali

# Movement Pruning per Transformer con Hugging Face + SparseML
# pip install transformers datasets sparseml

import torch
import torch.nn as nn
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from torch.optim import AdamW

# ===================================================================
# MOVEMENT PRUNING MANUALE (concetto base)
# ===================================================================

class MovementPruningLinear(nn.Module):
    """
    Layer Linear con movement pruning.
    Mantiene uno score per ogni peso: lo score viene ottimizzato
    durante il training. I pesi con score basso vengono pruned.
    """
    def __init__(self, in_features, out_features, pruning_ratio=0.5):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features) * 0.01)
        self.bias = nn.Parameter(torch.zeros(out_features))
        # Score inizializzati a zero: durante il training salgono per i pesi importanti
        self.scores = nn.Parameter(torch.zeros_like(self.weight))
        self.pruning_ratio = pruning_ratio
        self.mask = None

    def update_mask(self):
        """Aggiorna la maschera basandosi sugli score correnti."""
        k = int(self.scores.numel() * (1 - self.pruning_ratio))
        # Top-k scores: mantieni i pesi con score più alto
        threshold = torch.kthvalue(self.scores.flatten(), self.scores.numel() - k).values
        self.mask = (self.scores >= threshold).float().detach()

    def forward(self, x):
        # Applica la maschera durante il forward pass
        if self.mask is None:
            self.update_mask()
        masked_weight = self.weight * self.mask
        return nn.functional.linear(x, masked_weight, self.bias)


# ===================================================================
# PRUNING PRATICO CON TRANSFORMERS + torch.nn.utils.prune
# ===================================================================

def prune_transformer_attention_heads(model, heads_to_prune):
    """
    Pruna specifici attention heads da un modello BERT-like.
    heads_to_prune: dict {layer_idx: [head_idx_1, head_idx_2, ...]}
    """
    model.prune_heads(heads_to_prune)
    return model

# Esempio: pruning degli attention heads meno importanti
# Identificazione heads da pruning (basata su Taylor importance)
def compute_head_importance(model, dataloader, device):
    """
    Calcola l'importanza di ogni attention head usando Taylor expansion.
    Un head e importante se rimuoverlo aumenta molto la loss.
    """
    model.eval()
    head_importance = torch.zeros(
        model.config.num_hidden_layers,
        model.config.num_attention_heads
    ).to(device)

    for batch in dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch, output_attentions=True)
        loss = outputs.loss
        loss.backward()

        # Accumula gradienti per stimare l'importanza
        for layer_idx, layer in enumerate(model.bert.encoder.layer):
            attn = layer.attention.self
            # Importanza approssimata: |grad * weight| sommato per head
            grad = attn.value.weight.grad
            weight = attn.value.weight
            if grad is not None:
                importance = (grad * weight).abs().view(
                    model.config.num_attention_heads, -1
                ).sum(dim=-1)
                head_importance[layer_idx] += importance

    return head_importance


# ===================================================================
# STRUCTURED PRUNING DI ATTENTION HEADS CON BERT
# ===================================================================

model_name = "bert-base-uncased"
# model = AutoModelForSequenceClassification.from_pretrained(model_name)
# tokenizer = AutoTokenizer.from_pretrained(model_name)

# Strategia di pruning: rimuovi il 30% degli heads meno importanti
# Assumendo head_importance calcolata come sopra:
# heads_to_prune = {}
# n_heads_to_prune = int(0.3 * 12 * 12)  # 30% di 144 heads totali (12 layers x 12 heads)
# flat_importance = head_importance.flatten()
# _, indices = flat_importance.sort()
# for idx in indices[:n_heads_to_prune]:
#     layer_idx = idx.item() // 12
#     head_idx = idx.item() % 12
#     if layer_idx not in heads_to_prune:
#         heads_to_prune[layer_idx] = []
#     heads_to_prune[layer_idx].append(head_idx)
# pruned_model = prune_transformer_attention_heads(model, heads_to_prune)

print("Movement pruning e head importance pruning: schema implementato.")
print("Risultati tipici su BERT-base con 40% pruning attention:")
print("  - Speedup inferenza: 1.3-1.5x")
print("  - Dimensione modello: -35%")
print("  - Accuratezza GLUE: -0.5 a -1.5 punti")

import torch
import torch.nn as nn
import copy
from typing import Dict, List

# ===================================================================
# ITERATIVE MAGNITUDE PRUNING (Lottery Ticket Hypothesis)
# ===================================================================

def save_initial_weights(model: nn.Module) -> Dict[str, torch.Tensor]:
    """Salva i pesi iniziali del modello (prima del training)."""
    return {
        name: param.data.clone()
        for name, param in model.named_parameters()
        if 'weight' in name
    }

def apply_mask_and_reinit(
    model: nn.Module,
    initial_weights: Dict[str, torch.Tensor],
    masks: Dict[str, torch.Tensor]
) -> nn.Module:
    """
    Reimposta i pesi ai valori iniziali con le maschere di pruning applicate.
    Questo e il passo critico della LTH: reinizializzare (non random, ma ai valori originali).
    """
    with torch.no_grad():
        for name, param in model.named_parameters():
            if name in initial_weights and name in masks:
                param.data = initial_weights[name] * masks[name]
    return model

def compute_pruning_masks(
    model: nn.Module,
    pruning_ratio: float
) -> Dict[str, torch.Tensor]:
    """Calcola le maschere di pruning per magnitude (L1)."""
    masks = {}
    for name, param in model.named_parameters():
        if 'weight' in name and param.dim() > 1:
            # Soglia globale per layer
            threshold = torch.quantile(param.abs(), pruning_ratio)
            masks[name] = (param.abs() >= threshold).float()
    return masks

def iterative_magnitude_pruning(
    model: nn.Module,
    train_fn,
    eval_fn,
    n_rounds: int = 5,
    prune_per_round: float = 0.20,
    epochs_per_round: int = 10
):
    """
    Implementazione dell'Iterative Magnitude Pruning (LTH).

    Algoritmo:
    1. Salva i pesi iniziali (w0)
    2. Addestra per N epoche
    3. Pruna il P% dei pesi con magnitudine minore
    4. Reinizializza i pesi sopravvissuti a w0
    5. Ripeti dal passo 2
    """
    # Step 1: Salva i pesi iniziali
    initial_weights = save_initial_weights(model)
    masks = {name: torch.ones_like(param)
              for name, param in model.named_parameters()
              if 'weight' in name}

    cumulative_pruned = 0.0
    results = []

    for round_idx in range(n_rounds):
        print(f"\n--- Round IMP {round_idx + 1}/{n_rounds} ---")

        # Step 2: Addestra il modello (con le maschere correnti applicate)
        train_fn(model, epochs=epochs_per_round, masks=masks)

        # Step 3: Calcola nuove maschere di pruning
        effective_prune = 1 - (1 - prune_per_round) ** (round_idx + 1)
        new_masks = compute_pruning_masks(model, effective_prune)

        # Step 4: Reinizializza con pesi iniziali e nuove maschere
        model = apply_mask_and_reinit(model, initial_weights, new_masks)
        masks = new_masks

        # Valutazione
        accuracy = eval_fn(model)
        total_sparsity = sum(
            (m == 0).float().mean().item()
            for m in masks.values()
        ) / len(masks)

        results.append({
            'round': round_idx + 1,
            'accuracy': accuracy,
            'sparsity': total_sparsity
        })

        print(f"Accuratezza: {accuracy:.2%} | Sparsita: {total_sparsity:.1%}")

    return model, results


# Risultati tipici IMP su ResNet-20 / CIFAR-10:
# Round 1 (20% pruned):  91.8% accuracy (baseline: 91.9%)
# Round 2 (36% pruned):  91.7% accuracy
# Round 3 (49% pruned):  91.5% accuracy
# Round 4 (59% pruned):  91.2% accuracy
# Round 5 (67% pruned):  90.8% accuracy  <- "winning ticket"
# Round 8 (83% pruned):  89.1% accuracy  <- accuratezza inizia a degradare
# Round 10 (89% pruned): 87.3% accuracy  <- soglia tipica fine utilita

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune
from torch.optim.lr_scheduler import CosineAnnealingLR

# ===================================================================
# WORKFLOW PRUNING ITERATIVO COMPLETO
# ===================================================================

def get_global_sparsity(model: nn.Module) -> float:
    """Calcola la sparsita globale del modello."""
    total_params = 0
    zero_params = 0
    for name, param in model.named_parameters():
        if 'weight' in name:
            total_params += param.numel()
            zero_params += (param == 0).sum().item()
    return zero_params / total_params if total_params > 0 else 0.0

def iterative_pruning_with_finetuning(
    model: nn.Module,
    train_loader,
    val_loader,
    target_sparsity: float = 0.70,
    n_pruning_steps: int = 7,
    finetune_epochs_per_step: int = 3,
    lr: float = 1e-4,
    device: str = 'cuda'
):
    """
    Pruning iterativo con fine-tuning post-pruning.

    Strategia: aumenta la sparsita gradualmente usando una schedule
    cubica (più aggressiva all'inizio, più conservativa alla fine).
    """
    model = model.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-5)
    criterion = nn.CrossEntropyLoss()
    history = []

    # Schedule di sparsita cubica
    sparsity_schedule = [
        1 - (1 - target_sparsity * (step / n_pruning_steps) ** 3)
        for step in range(1, n_pruning_steps + 1)
    ]

    print(f"Schedule sparsita: {[f'{s:.1%}' for s in sparsity_schedule]}")

    for step_idx, target_sparsity_step in enumerate(sparsity_schedule):
        print(f"\n=== Step {step_idx + 1}/{n_pruning_steps} | Target sparsita: {target_sparsity_step:.1%} ===")

        # Raccoglie tutti i parametri weight del modello
        parameters_to_prune = [
            (module, 'weight')
            for name, module in model.named_modules()
            if isinstance(module, (nn.Linear, nn.Conv2d))
        ]

        # Pruning globale L1
        prune.global_unstructured(
            parameters_to_prune,
            pruning_method=prune.L1Unstructured,
            amount=target_sparsity_step
        )

        actual_sparsity = get_global_sparsity(model)
        print(f"Sparsita effettiva: {actual_sparsity:.1%}")

        # Fine-tuning post-pruning
        scheduler = CosineAnnealingLR(optimizer, T_max=finetune_epochs_per_step)
        for epoch in range(finetune_epochs_per_step):
            model.train()
            train_loss = 0.0
            for batch_x, batch_y in train_loader:
                batch_x, batch_y = batch_x.to(device), batch_y.to(device)
                optimizer.zero_grad()
                out = model(batch_x)
                loss = criterion(out, batch_y)
                loss.backward()
                optimizer.step()
                train_loss += loss.item()
            scheduler.step()

        # Valutazione
        model.eval()
        correct = total = 0
        with torch.no_grad():
            for batch_x, batch_y in val_loader:
                batch_x, batch_y = batch_x.to(device), batch_y.to(device)
                pred = model(batch_x).argmax(dim=1)
                correct += (pred == batch_y).sum().item()
                total += batch_y.size(0)

        val_acc = correct / total
        history.append({'step': step_idx+1, 'sparsity': actual_sparsity, 'val_acc': val_acc})
        print(f"Val accuracy: {val_acc:.2%}")

    # Consolida le maschere (rende il pruning permanente)
    for module, param_name in parameters_to_prune:
        try:
            prune.remove(module, param_name)
        except ValueError:
            pass  # Già rimosso

    return model, history

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# ===================================================================
# COMBINAZIONE PRUNING + QUANTIZZAZIONE
# ===================================================================

# --- Approccio 1: Pruning strutturato + Quantizzazione INT8 ---
# Pruna prima (rimuove strutture), poi quantizza il modello ridotto

def prune_and_quantize_pipeline(model_name: str, prune_ratio: float = 0.30):
    """
    Pipeline: carica modello -> pruning strutturato -> quantizzazione INT8.
    """
    # Step 1: Carica modello full precision
    from transformers import AutoModelForSequenceClassification
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name,
        torch_dtype=torch.float32
    )

    print(f"Parametri originali: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M")

    # Step 2: Pruning L1 non strutturato globale
    parameters_to_prune = [
        (module, 'weight')
        for name, module in model.named_modules()
        if isinstance(module, nn.Linear) and 'classifier' not in name
    ]
    prune.global_unstructured(
        parameters_to_prune,
        pruning_method=prune.L1Unstructured,
        amount=prune_ratio
    )
    # Consolida maschere
    for module, param_name in parameters_to_prune:
        prune.remove(module, param_name)

    # Conta parametri zero
    zero_params = sum(
        (param == 0).sum().item()
        for name, param in model.named_parameters()
        if 'weight' in name
    )
    total_params = sum(
        param.numel()
        for name, param in model.named_parameters()
        if 'weight' in name
    )
    print(f"Sparsita dopo pruning: {zero_params/total_params:.1%}")

    # Step 3: Quantizzazione dinamica INT8 del modello pruned
    model_quantized = torch.quantization.quantize_dynamic(
        model,
        {nn.Linear},       # Quantizza solo layer Linear
        dtype=torch.qint8
    )

    return model_quantized


# --- Approccio 2: QLoRA su modello pre-pruned ---
# Per LLM: usa modelli già pruned + quantizzazione NF4 per fine-tuning

config_nf4 = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True
)

# Molti modelli su HuggingFace Hub sono già pruned E quantizzati:
# es. "microsoft/phi-2" (2.7B), "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
# Questi sono modelli "distilled + pruned" during pretraining.


# --- Benchmark memoria: pruning + quantizzazione ---
compression_results = [
    {"metodo": "FP32 (baseline)",      "sparsita": "0%",  "precisione": "FP32", "size_mb": 440},
    {"metodo": "Pruning 50%",           "sparsita": "50%", "precisione": "FP32", "size_mb": 220},
    {"metodo": "Quantizzazione INT8",   "sparsita": "0%",  "precisione": "INT8", "size_mb": 110},
    {"metodo": "Pruning 50% + INT8",    "sparsita": "50%", "precisione": "INT8", "size_mb":  55},
    {"metodo": "Pruning 70% + INT4",    "sparsita": "70%", "precisione": "INT4", "size_mb":  33},
]

print(f"\n{'Metodo':<28} {'Sparsita':>10} {'Precisione':>12} {'Dimensione':>12}")
print("-" * 65)
for r in compression_results:
    print(f"{r['metodo']:<28} {r['sparsità]:>10} {r['precisione']:>12} {r['size_mb']:>10} MB")

# Output (modello BERT-base ~440MB in FP32):
# Metodo                       Sparsita   Precisione   Dimensione
# FP32 (baseline)                    0%         FP32       440 MB
# Pruning 50%                       50%         FP32       220 MB
# Quantizzazione INT8                0%         INT8       110 MB
# Pruning 50% + INT8                50%         INT8        55 MB
# Pruning 70% + INT4                70%         INT4        33 MB