Magnitude pruning

nif.pruning.PolynomialMagnitudePruning controls one selected kernel of an eager PyTorch layer. It updates its magnitude mask on scheduled training forwards and reapplies that mask before every forward, including evaluation. The bias is unchanged unless it is explicitly selected as the parameter.

class nif.pruning.PolynomialMagnitudePruning(layer, name='w', *, initial_sparsity=0.1, final_sparsity=0.9, begin_step=0, end_step=200000, power=3.0, frequency=4000)

Attach reference-convention pruning to one kernel, leaving its bias alone.

The training counter starts at zero and increments before each training forward. Refresh masks at eligible steps, then assign weight = mask before *every forward (including eval). Call on_epoch_end() before checkpointing: optimizer updates can transiently regrow masked coordinates. Refresh ranks current magnitudes, keeps max(round(N*(1-s)), 1) as a threshold index, and retains all threshold ties. Actual sparsity can therefore differ from s.

Model state_dict contains mask, threshold, counter and schedule configuration. To resume, attach a controller to the same kernel in the new model before loading its state_dict, and separately restore the optimizer state. strip() remasks and removes only pruning control state/hooks, retaining the original Parameter object and genuine sparse values for normal export. Eager dense real floating kernels are supported with one forward followed by its backward per optimizer update, as in tutorial 7. Multiple pending forward graphs before a combined backward with coordinate gradients are unsupported because masking modifies the kernel in place. Tracing/export follows strip().

property mask
property step
property refresh_count
on_epoch_end()

Clean regrowth without advancing the counter or refreshing the mask.

strip()

Keep sparse weights and Parameter identity; remove the pruning state.

Lifecycle

  1. Attach the controller before creating or restoring a pruning checkpoint.

  2. For each optimizer update, run one forward followed by its backward and step.

  3. Call on_epoch_end() before saving a checkpoint or evaluating final weights.

  4. Call strip() before copying, quantizing, tracing or exporting the inference model. This keeps the sparse parameter values and removes pruning hooks and control buffers.

Optimizer updates can regrow masked entries until the mask is reapplied. Threshold ties mean the observed sparsity can differ from the requested schedule. Multiple pending forward graphs followed by a combined backward are unsupported when coordinate gradients are involved because masking changes the parameter in place.

The parameter name depends on the layer: native torch.nn.Linear uses name="weight", while the tutorial’s HyperLinearForSIREN uses name="w". The controller’s mask, threshold, counter and schedule are part of the model’s state_dict. To resume, attach a controller to the same kernel before loading that state, and restore the optimizer state separately.

See the homepage’s network pruning and quantization example for a small float32 NIF workflow, and the native Tutorial 7 notebook for the full multiscale experiment and dynamic int8 CPU inference.