1import tensorflow as tf
2import tensorflow_model_optimization as tfmot
3
4# Build a model
5model = tf.keras.Sequential([
6    tf.keras.layers.Dense(512, activation='relu', input_shape=(784,)),
7    tf.keras.layers.Dense(256, activation='relu'),
8    tf.keras.layers.Dense(10, activation='softmax')
9])
10
11# Define pruning parameters
12pruning_params = {
13    'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
14        initial_sparsity=0.0,
15        final_sparsity=0.5,       # Remove 50% of weights
16        begin_step=1000,
17        end_step=5000
18    )
19}
20
21# Apply pruning to the entire model
22pruned_model = tfmot.sparsity.keras.prune_low_magnitude(model, **pruning_params)
23
24# Compile and train with pruning callbacks
25pruned_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
26
27pruned_model.fit(
28    X_train, y_train,
29    epochs=10,
30    validation_data=(X_val, y_val),
31    callbacks=[tfmot.sparsity.keras.UpdatePruningStep()]
32)
33
34# Strip pruning wrappers for deployment
35final_model = tfmot.sparsity.keras.strip_pruning(pruned_model)

1import torch
2import torch.nn as nn
3import torch.nn.utils.prune as prune
4
5class SimpleNet(nn.Module):
6    def __init__(self):
7        super().__init__()
8        self.fc1 = nn.Linear(784, 512)
9        self.fc2 = nn.Linear(512, 256)
10        self.fc3 = nn.Linear(256, 10)
11
12    def forward(self, x):
13        x = torch.relu(self.fc1(x))
14        x = torch.relu(self.fc2(x))
15        return self.fc3(x)
16
17model = SimpleNet()
18
19# Unstructured L1 pruning — remove 30% of weights with smallest magnitude
20prune.l1_unstructured(model.fc1, name='weight', amount=0.3)
21
22# Check sparsity
23weight = model.fc1.weight
24sparsity = (weight == 0).sum().item() / weight.numel()
25print(f"fc1 sparsity: {sparsity:.1%}")  # 30.0%
26
27# Structured pruning — remove entire output channels
28prune.ln_structured(model.fc2, name='weight', amount=0.2, n=2, dim=0)
29
30# Make pruning permanent (remove forward hooks)
31prune.remove(model.fc1, 'weight')

1parameters_to_prune = [
2    (model.fc1, 'weight'),
3    (model.fc2, 'weight'),
4    (model.fc3, 'weight'),
5]
6
7# Remove 40% of weights globally (smallest magnitudes across all layers)
8prune.global_unstructured(
9    parameters_to_prune,
10    pruning_method=prune.L1Unstructured,
11    amount=0.4
12)
13
14# Check per-layer sparsity
15for name, module in model.named_modules():
16    if isinstance(module, nn.Linear):
17        sparsity = (module.weight == 0).sum().item() / module.weight.numel()
18        print(f"{name} sparsity: {sparsity:.1%}")

1def iterative_pruning(model, train_loader, val_loader, target_sparsity=0.9, rounds=5):
2    sparsity_per_round = 1 - (1 - target_sparsity) ** (1 / rounds)
3
4    for round_num in range(rounds):
5        # Prune
6        for name, module in model.named_modules():
7            if isinstance(module, nn.Linear):
8                prune.l1_unstructured(module, name='weight', amount=sparsity_per_round)
9
10        # Fine-tune for a few epochs
11        optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
12        for epoch in range(3):
13            train_one_epoch(model, train_loader, optimizer)
14
15        val_acc = evaluate(model, val_loader)
16        current_sparsity = get_model_sparsity(model)
17        print(f"Round {round_num+1}: sparsity={current_sparsity:.1%}, val_acc={val_acc:.4f}")
18
19    # Make pruning permanent
20    for name, module in model.named_modules():
21        if isinstance(module, nn.Linear):
22            prune.remove(module, 'weight')
23
24    return model

1# 1. Train the full model
2model = SimpleNet()
3original_weights = {name: param.clone() for name, param in model.named_parameters()}
4train(model, train_loader, epochs=20)
5
6# 2. Prune to find the mask
7for name, module in model.named_modules():
8    if isinstance(module, nn.Linear):
9        prune.l1_unstructured(module, name='weight', amount=0.8)
10
11masks = {name: module.weight_mask.clone() for name, module in model.named_modules()
12         if hasattr(module, 'weight_mask')}
13
14# 3. Reset to original weights and apply mask
15for name, param in model.named_parameters():
16    if name in original_weights:
17        param.data = original_weights[name] * masks.get(name.replace('.weight', '.weight_mask'), 1)
18
19# 4. Retrain the sparse subnetwork
20train(model, train_loader, epochs=20)

1def get_model_sparsity(model):
2    total_params = 0
3    zero_params = 0
4    for param in model.parameters():
5        total_params += param.numel()
6        zero_params += (param == 0).sum().item()
7    return zero_params / total_params
8
9def get_model_size_mb(model):
10    param_size = sum(p.nelement() * p.element_size() for p in model.parameters())
11    return param_size / (1024 * 1024)
12
13print(f"Sparsity: {get_model_sparsity(model):.1%}")
14print(f"Size: {get_model_size_mb(model):.2f} MB")