Pruning Basics › Code Implementation
Language AI HandbookAll chapters Pruning Basics › Code ImplementationOne section · Sections in this chapter7 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us implement the key components of a pruning pipeline using PyTorch's built-in pruning utilities. We will apply magnitude pruning to a small feedforward network and examine the sparsity-accuracy tradeoff as we increase the pruning ratio.
Setup and Network Definition
We start by importing the libraries and defining a small feedforward network for text classification.
# Reproducibility
torch.manual_seed(42)
np.random.seed(42)
# Small feedforward network for binary text classification
class TextClassifier(nn.Module):
def __init__(self, input_dim=100, hidden_dim=256, output_dim=2):
super().__init__()
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.relu = nn.ReLU()
self.dropout = nn.Dropout(0.3)
self.fc2 = nn.Linear(hidden_dim, hidden_dim // 2)
self.fc3 = nn.Linear(hidden_dim // 2, output_dim)
def forward(self, x):
x = self.dropout(self.relu(self.fc1(x)))
x = self.relu(self.fc2(x))
return self.fc3(x)
model = TextClassifier(input_dim=100, hidden_dim=256, output_dim=2)Total parameters: 59,010 fc1.weight: torch.Size([256, 100]) -> 25,600 params fc1.bias: torch.Size([256]) -> 256 params fc2.weight: torch.Size([128, 256]) -> 32,768 params fc2.bias: torch.Size([128]) -> 128 params fc3.weight: torch.Size([2, 128]) -> 256 params fc3.bias: torch.Size([2]) -> 2 params
The model has three fully connected layers. The first layer maps 100-dimensional TF-IDF-style features to 256 hidden units, the second halves that to 128, and the third produces the two-class logits. With three layers of sizes 100x256, 256x128, and 128x2, the weight matrices contain 25,600, 32,768, and 256 parameters respectively. The first two layers together account for over 98% of the model's parameters and are the primary targets for pruning.
Synthetic Data and Training
We generate synthetic TF-IDF-like features and train the model to convergence before pruning. The synthetic task is simple: predict which of two word groups has a higher aggregate frequency in the input. This gives us a well-defined signal-to-noise ratio that makes the effect of pruning easy to interpret.
# Synthetic TF-IDF-like data: sparse non-negative features
N_TRAIN, N_TEST = 2000, 500
INPUT_DIM = 100
X_train = torch.relu(torch.randn(N_TRAIN, INPUT_DIM))
y_train = (X_train[:, :10].sum(dim=1) > X_train[:, 10:20].sum(dim=1)).long()
X_test = torch.relu(torch.randn(N_TEST, INPUT_DIM))
y_test = (X_test[:, :10].sum(dim=1) > X_test[:, 10:20].sum(dim=1)).long()
train_dataset = TensorDataset(X_train, y_train)
test_dataset = TensorDataset(X_test, y_test)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)def train_model(model, train_loader, epochs=20, lr=1e-3):
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()
model.train()
losses = []
for epoch in range(epochs):
epoch_loss = 0.0
for X_batch, y_batch in train_loader:
optimizer.zero_grad()
logits = model(X_batch)
loss = criterion(logits, y_batch)
loss.backward()
optimizer.step()
epoch_loss += loss.item() * X_batch.size(0)
losses.append(epoch_loss / len(train_loader.dataset))
return losses
def evaluate_model(model, data_loader):
model.eval()
correct = 0
total = 0
with torch.no_grad():
for X_batch, y_batch in data_loader:
preds = model(X_batch).argmax(dim=1)
correct += (preds == y_batch).sum().item()
total += y_batch.size(0)
return correct / total
training_losses = train_model(model, train_loader, epochs=20)
baseline_accuracy = evaluate_model(model, test_loader)Baseline test accuracy: 0.962 Final training loss: 0.0160
Applying Magnitude Pruning
PyTorch's torch.nn.utils.prune module provides a clean API for applying pruning masks. We apply global unstructured magnitude pruning, which considers all weights across all specified layers together and prunes the globally smallest weights regardless of which layer they are in.
The global_unstructured function assigns a pruning mask to each parameter: a binary tensor of the same shape as the weight matrix, where 1 means "keep" and 0 means "prune." The actual weight values are not immediately zeroed; instead, PyTorch registers a forward hook that applies the mask on each forward pass. The prune.remove call makes the pruning permanent by baking the zeros directly into the weight tensor and removing the hook.
def apply_global_magnitude_pruning(model, sparsity, make_permanent=True):
pruned_model = copy.deepcopy(model)
parameters_to_prune = [
(pruned_model.fc1, "weight"),
(pruned_model.fc2, "weight"),
(pruned_model.fc3, "weight"),
]
prune.global_unstructured(
parameters_to_prune,
pruning_method=prune.L1Unstructured,
amount=sparsity,
)
if make_permanent:
# Remove the reparameterization while preserving the zeroed weights.
for module, name in parameters_to_prune:
prune.remove(module, name)
return pruned_model
def make_pruning_permanent(model):
"""Bake active pruning masks into each linear layer's weight tensor."""
for module in model.modules():
if isinstance(module, nn.Linear) and hasattr(module, "weight_mask"):
prune.remove(module, "weight")
def measure_sparsity(model):
total_weights = 0
zero_weights = 0
for module in model.modules():
if isinstance(module, nn.Linear):
total_weights += module.weight.numel()
zero_weights += (module.weight == 0).sum().item()
return zero_weights / total_weights if total_weights > 0 else 0.0# Evaluate performance across a range of sparsity levels (no fine-tuning)
sparsity_levels = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.95]
results = []
for target_sparsity in sparsity_levels:
if target_sparsity == 0.0:
pruned = copy.deepcopy(model)
else:
pruned = apply_global_magnitude_pruning(model, target_sparsity)
actual_sparsity = measure_sparsity(pruned)
accuracy = evaluate_model(pruned, test_loader)
results.append(
{
"target_sparsity": target_sparsity,
"actual_sparsity": actual_sparsity,
"accuracy": accuracy,
}
) Sparsity Actual Sparsity Test Accuracy
----------------------------------------------
0% 0.000 0.962
10% 0.100 0.966
20% 0.200 0.964
30% 0.300 0.962
40% 0.400 0.960
50% 0.500 0.954
60% 0.600 0.958
70% 0.700 0.936
80% 0.800 0.952
90% 0.900 0.900
95% 0.950 0.822Notice how accuracy holds up well through moderate sparsity but degrades more sharply past 80 to 90%. This pattern is characteristic of magnitude pruning without any fine-tuning. The model tolerates removal of small weights easily, because those weights were contributing little to begin with. But eventually the weights that are "small" by magnitude are still doing meaningful work, and removing them produces noticeable accuracy drops. The sparsity-accuracy curve is not linear; it has a long flat region followed by a steep cliff.
Fine-Tuning After Pruning
Fine-tuning the pruned model allows it to redistribute the workload among the remaining weights and recover lost accuracy. The surviving weights adjust their magnitudes to compensate for the information that was lost with the pruned weights. This is possible because neural networks are highly overdetermined: multiple weight configurations can encode the same function, and fine-tuning finds a new configuration within the constraints imposed by the sparse mask.
# Prune to 70% and fine-tune
TARGET_SPARSITY = 0.70
pruned_finetuned = apply_global_magnitude_pruning(
model, TARGET_SPARSITY, make_permanent=False
)
accuracy_before_ft = evaluate_model(pruned_finetuned, test_loader)
# Fine-tune with a reduced learning rate while the pruning masks stay active.
ft_losses = train_model(pruned_finetuned, train_loader, epochs=10, lr=5e-4)
make_pruning_permanent(pruned_finetuned)
accuracy_after_ft = evaluate_model(pruned_finetuned, test_loader)
actual_sparsity_ft = measure_sparsity(pruned_finetuned)Target sparsity: 70% Actual sparsity: 0.700 Accuracy before fine-tuning: 0.936 Accuracy after fine-tuning: 0.972 Accuracy recovery: +0.036
Fine-tuning even for a modest number of steps substantially recovers the accuracy lost to pruning. The pruned model learns to compensate for the missing weights by adjusting the values of surviving connections. This confirms the practical importance of including a fine-tuning phase after pruning: without it, the accuracy-sparsity curve clips off early; with it, you can achieve much higher sparsity at the same accuracy level.
The weight magnitude distribution provides a useful window into what pruning does to the parameter distribution. Before pruning, weights from a converged model typically form a roughly bell-shaped distribution centered near zero with long tails. After pruning, the near-zero region is surgically emptied, leaving only the weights that survived the magnitude threshold. The remaining weights are those in the long tails of the original distribution.
Gradual Pruning Schedule
Now we implement a simplified gradual pruning schedule: instead of pruning everything at once, we prune incrementally during training. The model is re-initialized with random weights for this experiment so we can observe the full gradual pruning trajectory from scratch.
def gradual_pruning_schedule(final_sparsity, n_steps, current_step):
if current_step >= n_steps:
return final_sparsity
fraction = current_step / n_steps
return final_sparsity * (1 - (1 - fraction) ** 3)
def train_with_gradual_pruning(
model_init,
train_loader,
final_sparsity=0.70,
n_prune_steps=10,
total_epochs=30,
lr=1e-3,
):
model_gp = copy.deepcopy(model_init)
for module in model_gp.modules():
if isinstance(module, nn.Linear):
nn.init.xavier_uniform_(module.weight)
nn.init.zeros_(module.bias)
optimizer = torch.optim.Adam(model_gp.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()
prune_interval = max(1, total_epochs // n_prune_steps)
accuracies = []
sparsities = []
for epoch in range(1, total_epochs + 1):
model_gp.train()
for X_batch, y_batch in train_loader:
optimizer.zero_grad()
loss = criterion(model_gp(X_batch), y_batch)
loss.backward()
optimizer.step()
prune_step = epoch // prune_interval
target_s = gradual_pruning_schedule(
final_sparsity, n_prune_steps, prune_step
)
if target_s > 0 and epoch % prune_interval == 0:
params = [
(model_gp.fc1, "weight"),
(model_gp.fc2, "weight"),
(model_gp.fc3, "weight"),
]
current_s = measure_sparsity(model_gp)
incremental_amount = (target_s - current_s) / (1 - current_s)
if incremental_amount > 0:
prune.global_unstructured(
params,
pruning_method=prune.L1Unstructured,
amount=incremental_amount,
)
acc = evaluate_model(model_gp, test_loader)
sp = measure_sparsity(model_gp)
accuracies.append(acc)
sparsities.append(sp)
make_pruning_permanent(model_gp)
return model_gp, accuracies, sparsities
gradual_model, gp_accuracies, gp_sparsities = train_with_gradual_pruning(
model, train_loader, final_sparsity=0.70, n_prune_steps=10, total_epochs=30
)
final_gp_accuracy = evaluate_model(gradual_model, test_loader)
final_gp_sparsity = measure_sparsity(gradual_model)Gradual pruning - final sparsity: 0.700 Gradual pruning - final accuracy: 0.958 Baseline accuracy: 0.962 Accuracy gap from baseline: 0.004
The gradual pruning trajectory reveals a characteristic pattern: accuracy dips slightly each time a new pruning step is applied, then recovers as training continues before the next pruning step. The dips are small because each individual pruning step removes only a small fraction of the total weight budget, and the model adapts quickly.
Key Parameters
The key parameters for weight pruning are:
amount(sparsity): The fraction of weights to prune (for example, 0.7 for 70% sparsity). Higher values compress more aggressively but risk greater accuracy loss. The optimal value depends on the model architecture, task, and available fine-tuning compute.pruning_method: The criterion for selecting which weights to remove.L1Unstructuredapplies magnitude pruning;RandomUnstructuredprovides a random baseline; custom methods can implement gradient-based or activation-based criteria.global_unstructuredvs. per-layer: Global pruning considers all layers together and allocates sparsity based on magnitude across the entire network, naturally pruning sensitive layers less. Per-layer pruning applies a uniform sparsity ratio to each layer independently, which may under-prune critical layers and over-prune redundant ones.- Fine-tuning learning rate: After pruning, fine-tuning with a reduced learning rate (typically 10 to 50% of the original) prevents large weight updates that could disrupt the remaining sparse structure.
- Gradual pruning interval: How often sparsity increases during gradual pruning. More frequent updates create a smoother sparsity trajectory; less frequent updates allow more recovery time between pruning steps.