---
title: ModelPruning
description: >-
  Model pruning Callback, using PyTorch’s prune utilities. This callback is
  responsible of pruning networks

  parameters during training.
canonical: >-
  https://lightning.ai/docs/pytorch/stable/optional-api/api/lightning.pytorch.callbacks.ModelPruning
---
# ModelPruning {#modelpruning}

### `class lightning.pytorch.callbacks.ModelPruning(pruning_fn, parameters_to_prune=(), parameter_names=None, use_global_unstructured=True, amount=0.5, apply_pruning=True, make_pruning_permanent=True, use_lottery_ticket_hypothesis=True, resample_parameters=False, pruning_dim=None, pruning_norm=None, verbose=0, prune_on_train_epoch_end=True)` {#lightning.pytorch.callbacks.ModelPruning}

Bases: [`Callback`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.callbacks.Callback.md#lightning.pytorch.callbacks.Callback "lightning.pytorch.callbacks.callback.Callback")

Model pruning Callback, using PyTorch’s prune utilities. This callback is responsible of pruning networks parameters during training.

To learn more about pruning with PyTorch, please take a look at [this tutorial](https://pytorch.org/tutorials/intermediate/pruning_tutorial.html).

Warning

This is an [experimental](/docs/pytorch/2.6.6/reference/versioning.md#experimental-api) feature.

```python
parameters_to_prune = [(model.mlp_1, "weight"), (model.mlp_2, "weight")]

trainer = Trainer(
    callbacks=[
        ModelPruning(
            pruning_fn="l1_unstructured",
            parameters_to_prune=parameters_to_prune,
            amount=0.01,
            use_global_unstructured=True,
        )
    ]
)
```

When `parameters_to_prune` is `None`, `parameters_to_prune` will contain all parameters from the model. The user can override `filter_parameters_to_prune` to filter any `nn.Module` to be pruned.

Parameters:

*   **pruning\_fn** (`Union`\[[`Callable`](https://docs.python.org/3/library/typing.html#typing.Callable "(in Python v3.15)"), [`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")\]) – Function from torch.nn.utils.prune module or your own PyTorch `BasePruningMethod` subclass. Can also be string e.g. “l1\_unstructured”. See pytorch docs for more details.
    
*   **parameters\_to\_prune** ([`Sequence`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence "(in Python v3.15)")\[[`tuple`](https://docs.python.org/3/builtins/stdtypes.html#tuple "(in Python v3.15)")\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")\]\]) – List of tuples `(nn.Module, "parameter_name_string")`.
    
*   **parameter\_names** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`list`](https://docs.python.org/3/builtins/stdtypes.html#list "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")\]\]) – List of parameter names to be pruned from the nn.Module. Can either be `"weight"` or `"bias"`.
    
*   **use\_global\_unstructured** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Whether to apply pruning globally on the model. If `parameters_to_prune` is provided, global unstructured will be restricted on them.
    
*   **amount** (`Union`\[[`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)"), [`float`](https://docs.python.org/3/builtins/functions.html#float "(in Python v3.15)"), [`Callable`](https://docs.python.org/3/library/typing.html#typing.Callable "(in Python v3.15)")\[\[[`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")\], `Union`\[[`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)"), [`float`](https://docs.python.org/3/builtins/functions.html#float "(in Python v3.15)")\]\]\]) –
    
    Quantity of parameters to prune:
    
    *   `float`. Between 0.0 and 1.0. Represents the fraction of parameters to prune.
        
    *   `int`. Represents the absolute number of parameters to prune.
        
    *   `Callable`. For dynamic values. Will be called every epoch. Should return a value.
        
    
*   **apply\_pruning** (`Union`\[[`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)"), [`Callable`](https://docs.python.org/3/library/typing.html#typing.Callable "(in Python v3.15)")\[\[[`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")\], [`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")\]\]) –
    
    Whether to apply pruning.
    
    *   `bool`. Always apply it or not.
        
    *   `Callable[[epoch], bool]`. For dynamic values. Will be called every epoch.
        
    
*   **make\_pruning\_permanent** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Whether to remove all reparameterization pre-hooks and apply masks when training ends or the model is saved.
    
*   **use\_lottery\_ticket\_hypothesis** (`Union`\[[`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)"), [`Callable`](https://docs.python.org/3/library/typing.html#typing.Callable "(in Python v3.15)")\[\[[`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")\], [`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")\]\]) –
    
    See [The lottery ticket hypothesis](https://arxiv.org/abs/1803.03635):
    
    *   `bool`. Whether to apply it or not.
        
    *   `Callable[[epoch], bool]`. For dynamic values. Will be called every epoch.
        
    
*   **resample\_parameters** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – Used with `use_lottery_ticket_hypothesis`. If True, the model parameters will be resampled, otherwise, the exact original parameters will be used.
    
*   **pruning\_dim** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")\]) – If you are using a structured pruning method you need to specify the dimension.
    
*   **pruning\_norm** ([`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")\]) – If you are using `ln_structured` you need to specify the norm.
    
*   **verbose** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Verbosity level. 0 to disable, 1 to log overall sparsity, 2 to log per-layer sparsity
    
*   **prune\_on\_train\_epoch\_end** ([`bool`](https://docs.python.org/3/builtins/functions.html#bool "(in Python v3.15)")) – whether to apply pruning at the end of the training epoch. If this is `False`, then the check runs at the end of the validation epoch.
    

Raises:

**MisconfigurationException** – If `parameter_names` is neither `"weight"` nor `"bias"`, if the provided `pruning_fn` is not supported, if `pruning_dim` is not provided when `"unstructured"`, if `pruning_norm` is not provided when `"ln_structured"`, if `pruning_fn` is neither `str` nor [`torch.nn.utils.prune.BasePruningMethod`](https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.prune.BasePruningMethod.html#torch.nn.utils.prune.BasePruningMethod "(in PyTorch v2.15)"), or if `amount` is none of `int`, `float` and `Callable`.

#### `apply_lottery_ticket_hypothesis()` {#lightning.pytorch.callbacks.ModelPruning.apply_lottery_ticket_hypothesis}

Lottery ticket hypothesis algorithm (see page 2 of the paper): :rtype: [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

1.  Randomly initialize a neural network $`f(x; \theta_0)`$ (where $`\theta_0 \sim \mathcal{D}_\theta`$).
    
2.  Train the network for $`j`$ iterations, arriving at parameters $`\theta_j`$.
    
3.  Prune $`p\%`$ of the parameters in $`\theta_j`$, creating a mask $`m`$.
    
4.  Reset the remaining parameters to their values in $`\theta_0`$, creating the winning ticket $`f(x; m \odot \theta_0)`$.
    

This function implements the step 4.

The `resample_parameters` argument can be used to reset the parameters with a new $`\theta_z \sim \mathcal{D}_\theta`$

#### `apply_pruning(amount)` {#lightning.pytorch.callbacks.ModelPruning.apply_pruning}

Applies pruning to `parameters_to_prune`.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `filter_parameters_to_prune(parameters_to_prune=())` {#lightning.pytorch.callbacks.ModelPruning.filter_parameters_to_prune}

This function can be overridden to control which module to prune.

Return type:

[`Sequence`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence "(in Python v3.15)")\[[`tuple`](https://docs.python.org/3/builtins/stdtypes.html#tuple "(in Python v3.15)")\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")\]\]

#### `make_pruning_permanent(module)` {#lightning.pytorch.callbacks.ModelPruning.make_pruning_permanent}

Removes pruning buffers from any pruned modules.

Adapted from [https://github.com/pytorch/pytorch/blob/v1.7.1/torch/nn/utils/prune.py#L1118-L1122](https://github.com/pytorch/pytorch/blob/v1.7.1/torch/nn/utils/prune.py#L1118-L1122)

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_save_checkpoint(trainer, pl_module, checkpoint)` {#lightning.pytorch.callbacks.ModelPruning.on_save_checkpoint}

Called when saving a checkpoint to give you a chance to store anything else you might want to save.

Parameters:

*   **trainer** ([`Trainer`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md#lightning.pytorch.trainer.trainer.Trainer "lightning.pytorch.trainer.trainer.Trainer")) – the current [`Trainer`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.trainer.trainer.Trainer.md#lightning.pytorch.trainer.trainer.Trainer "lightning.pytorch.trainer.trainer.Trainer") instance.
    
*   **pl\_module** ([`LightningModule`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.core.LightningModule.md#lightning.pytorch.core.LightningModule "lightning.pytorch.core.module.LightningModule")) – the current [`LightningModule`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.core.LightningModule.md#lightning.pytorch.core.LightningModule "lightning.pytorch.core.LightningModule") instance.
    
*   **checkpoint** ([`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – the checkpoint dictionary that will be saved.
    

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_train_end(trainer, pl_module)` {#lightning.pytorch.callbacks.ModelPruning.on_train_end}

Called when the train ends.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_train_epoch_end(trainer, pl_module)` {#lightning.pytorch.callbacks.ModelPruning.on_train_epoch_end}

Called when the train epoch ends.

To access all batch outputs at the end of the epoch, you can cache step outputs as an attribute of the [`lightning.pytorch.core.LightningModule`](/docs/pytorch/2.6.6/optional-api/api/lightning.pytorch.core.LightningModule.md#lightning.pytorch.core.LightningModule "lightning.pytorch.core.LightningModule") and access them in this hook:

```python
class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []

    def training_step(self):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss


class MyCallback(L.Callback):
    def on_train_epoch_end(self, trainer, pl_module):
        # do something with all training_step outputs, for example:
        epoch_mean = torch.stack(pl_module.training_step_outputs).mean()
        pl_module.log("training_epoch_mean", epoch_mean)
        # free up the memory
        pl_module.training_step_outputs.clear()
```

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_validation_epoch_end(trainer, pl_module)` {#lightning.pytorch.callbacks.ModelPruning.on_validation_epoch_end}

Called when the val epoch ends.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `static sanitize_parameters_to_prune(pl_module, parameters_to_prune=(), parameter_names=())` {#lightning.pytorch.callbacks.ModelPruning.sanitize_parameters_to_prune}

This function is responsible of sanitizing `parameters_to_prune` and `parameter_names`. If `parameters_to_prune is None`, it will be generated with all parameters of the model.

Raises:

**MisconfigurationException** – If `parameters_to_prune` doesn’t exist in the model, or if `parameters_to_prune` is neither a list nor a tuple.

Return type:

[`Sequence`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence "(in Python v3.15)")\[[`tuple`](https://docs.python.org/3/builtins/stdtypes.html#tuple "(in Python v3.15)")\[[`Module`](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "(in PyTorch v2.15)"), [`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")\]\]

#### `setup(trainer, pl_module, stage)` {#lightning.pytorch.callbacks.ModelPruning.setup}

Called when fit, validate, test, predict, or tune begins.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

---

**Navigation**

- Parent: [Api](/docs/pytorch/2.6.6/optional-api/api.md)