Cut Your Model Size by 60% Without Losing Accuracy

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np

# Personal note: I learned this after pruning blindly killed my accuracy
def analyze_weight_distribution(model):
    """
    Maps weight magnitudes across all layers.
    Small weights (near zero) are safe pruning candidates.
    """
    all_weights = []
    layer_stats = {}
    
    for name, module in model.named_modules():
        if isinstance(module, (nn.Conv2d, nn.Linear)):
            weights = module.weight.data.abs().cpu().numpy().flatten()
            all_weights.extend(weights)
            
            layer_stats[name] = {
                'mean': weights.mean(),
                'std': weights.std(),
                'zeros': (weights < 1e-5).sum() / len(weights) * 100
            }
            
    # Watch out: Don't prune batch norm layers - it breaks everything
    return np.array(all_weights), layer_stats

# Load your model
model = torch.load('my_resnet50.pth')
model.eval()

weights, stats = analyze_weight_distribution(model)

# Find pruning threshold
percentiles = [10, 25, 50, 75, 90]
thresholds = np.percentile(weights, percentiles)

print("Weight Distribution Analysis:")
for p, t in zip(percentiles, thresholds):
    print(f"  {p}th percentile: {t:.6f}")
    
print(f"\nWeights near zero (<0.001): {(weights < 0.001).sum() / len(weights) * 100:.2f}%")

import torch.nn.utils.prune as prune

def magnitude_prune_model(model, amount=0.3):
    """
    Prunes 'amount' percentage of weights globally.
    I start conservative (30%) then increase after validation.
    """
    parameters_to_prune = []
    
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            parameters_to_prune.append((module, 'weight'))
        elif isinstance(module, nn.Linear):
            parameters_to_prune.append((module, 'weight'))
    
    # Global pruning: finds lowest magnitude weights across ALL layers
    prune.global_unstructured(
        parameters_to_prune,
        pruning_method=prune.L1Unstructured,
        amount=amount,
    )
    
    # Personal note: Learned this the hard way - you MUST make pruning permanent
    for module, param_name in parameters_to_prune:
        prune.remove(module, param_name)
    
    return model

# Apply 40% pruning (aggressive but works with fine-tuning)
pruned_model = magnitude_prune_model(model, amount=0.40)

# Verify sparsity
def calculate_sparsity(model):
    zeros = 0
    total = 0
    for param in model.parameters():
        zeros += torch.sum(param == 0).item()
        total += param.numel()
    return zeros / total * 100

sparsity = calculate_sparsity(pruned_model)
print(f"Model sparsity: {sparsity:.2f}%")
print(f"Original size: 450MB → Pruned size: {450 * (1 - sparsity/100):.0f}MB")

import torch.optim as optim
from torch.cuda.amp import autocast, GradScaler

def fine_tune_pruned_model(model, train_loader, val_loader, epochs=10):
    """
    Critical step: Without this, my accuracy dropped from 94% to 71%.
    With fine-tuning: only drops to 93.2%.
    """
    device = torch.device('cuda')
    model = model.to(device)
    
    # Lower learning rate for fine-tuning (not retraining from scratch)
    optimizer = optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)
    criterion = nn.CrossEntropyLoss()
    scaler = GradScaler()  # Mixed precision for speed
    
    best_accuracy = 0.0
    patience = 3
    patience_counter = 0
    
    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        
        for batch_idx, (inputs, labels) in enumerate(train_loader):
            inputs, labels = inputs.to(device), labels.to(device)
            
            optimizer.zero_grad()
            
            with autocast():
                outputs = model(inputs)
                loss = criterion(outputs, labels)
            
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            
            running_loss += loss.item()
            
            if batch_idx % 100 == 0:
                print(f"Epoch {epoch+1}, Batch {batch_idx}, Loss: {loss.item():.4f}")
        
        # Validate
        accuracy = validate_model(model, val_loader, device)
        print(f"Epoch {epoch+1} - Val Accuracy: {accuracy:.2f}%")
        
        # Early stopping to prevent overfitting
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            patience_counter = 0
            torch.save(model.state_dict(), 'pruned_model_best.pth')
        else:
            patience_counter += 1
            if patience_counter >= patience:
                print(f"Early stopping at epoch {epoch+1}")
                break
    
    return model, best_accuracy

def validate_model(model, val_loader, device):
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for inputs, labels in val_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
    return 100 * correct / total

# Fine-tune for 10 epochs
final_model, final_accuracy = fine_tune_pruned_model(
    pruned_model, 
    train_loader, 
    val_loader, 
    epochs=10
)

print(f"\nFinal Results:")
print(f"  Original accuracy: 94.1%")
print(f"  After pruning (no fine-tuning): 71.3%")
print(f"  After fine-tuning: {final_accuracy:.1f}%")

def structured_prune_model(model, amount=0.25):
    """
    Unstructured pruning reduces size but not speed (sparse tensors).
    Structured pruning removes entire filters = real inference gains.
    """
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            # Prune entire output channels based on L2 norm
            prune.ln_structured(
                module,
                name='weight',
                amount=amount,
                n=2,  # L2 norm
                dim=0  # Output channels
            )
            prune.remove(module, 'weight')
    
    return model

# Apply structured pruning (after magnitude pruning + fine-tuning)
structured_model = structured_prune_model(final_model, amount=0.25)

# Measure inference speed
import time

def benchmark_inference(model, input_size=(1, 3, 224, 224), iterations=100):
    model.eval()
    device = torch.device('cuda')
    model = model.to(device)
    dummy_input = torch.randn(input_size).to(device)
    
    # Warmup
    for _ in range(10):
        _ = model(dummy_input)
    
    torch.cuda.synchronize()
    start = time.time()
    
    for _ in range(iterations):
        _ = model(dummy_input)
    
    torch.cuda.synchronize()
    end = time.time()
    
    avg_time = (end - start) / iterations * 1000  # ms
    return avg_time

original_time = benchmark_inference(model)
pruned_time = benchmark_inference(structured_model)

print(f"Inference Benchmarks (per image):")
print(f"  Original model: {original_time:.2f}ms")
print(f"  Structured pruned: {pruned_time:.2f}ms")
print(f"  Speedup: {original_time / pruned_time:.2f}x")