Skip to content

PigeonColumba compressaTheme

Lesson 6 of 10

Global magnitude pruning

Magnitude pruning keeps the largest weights and zeros the rest. Global pruning ranks the weights of all layers together, so a layer full of tiny weights can lose most of them while a sensitive layer keeps more.

Implement global_masks(weights, sparsity): given a list of weight matrices and the fraction to remove, return one 0/1 mask per matrix (same shapes) that removes exactly round(sparsity * total) weights, the smallest in absolute value across all matrices.

Then the training loop multiplies weights by their mask after every step so pruned weights stay dead (lab/nn.py:MLP.apply_masks).

Your code

import numpy as np


def global_masks(weights, sparsity):
    """Return a list of 0/1 float arrays, one per weight matrix."""
    raise NotImplementedError

Checking your code and saving a draft need JavaScript, which is turned off in this browser. Turn it on and reload the page to check your code. You can still edit it here, and the hints and the reference solution are below.

Pigeon's hints (2)

Hint 1 of 2

Concatenate abs(w).ravel() for every matrix, rank once, then split back by sizes.

Hint 2 of 2

np.argsort(...)[:k] gives the indices of the k smallest.

The reference solution

import numpy as np


def global_masks(weights, sparsity):
    flat = np.concatenate([np.abs(w).ravel() for w in weights])
    k = int(round(sparsity * flat.size))
    drop = np.zeros(flat.size, bool)
    drop[np.argsort(flat, kind="stable")[:k]] = True
    masks, i = [], 0
    for w in weights:
        masks.append((~drop[i:i + w.size]).reshape(w.shape).astype(float))
        i += w.size
    return masks

Where this is used