docstr: 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 :ref:`experimental <versioning:Experimental API>` feature.
  
  .. code-block:: 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.
  
  Args:
  
      pruning_fn: 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: List of tuples ``(nn.Module, "parameter_name_string")``.
  
      parameter_names: List of parameter names to be pruned from the nn.Module.
          Can either be ``"weight"`` or ``"bias"``.
  
      use_global_unstructured: Whether to apply pruning globally on the model.
          If ``parameters_to_prune`` is provided, global unstructured will be restricted on them.
  
      amount: 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: Whether to apply pruning.
  
          - ``bool``. Always apply it or not.
          - ``Callable[[epoch], bool]``. For dynamic values. Will be called every epoch.
  
      make_pruning_permanent: Whether to remove all reparametrization pre-hooks and apply masks
          when training ends or the model is saved.
  
      use_lottery_ticket_hypothesis: 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: Used with ``use_lottery_ticket_hypothesis``. If True, the model parameters will
          be resampled, otherwise, the exact original parameters will be used.
  
      pruning_dim: If you are using a structured pruning method you need to specify the dimension.
  
      pruning_norm: If you are using ``ln_structured`` you need to specify the norm.
  
      verbose: Verbosity level. 0 to disable, 1 to log overall sparsity, 2 to log per-layer sparsity
  
      prune_on_train_epoch_end: 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 :class:`torch.nn.utils.prune.BasePruningMethod`, or
          if ``amount`` is none of ``int``, ``float`` and ``Callable``.
docstr_tok: ["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",":","ref",":","`","experimental","<","versioning",":","Experimental","API",">","`","feature",".","..","code-block",":",":","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",".","Args",":","pruning_fn",":","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",":","List","of","tuples","``","(","nn.Module",",","``","parameter_name_string","''",")","``",".","parameter_names",":","List","of","parameter","names","to","be","pruned","from","the","nn.Module",".","Can","either","be","``","``","weight","''","``","or","``","``","bias","''","``",".","use_global_unstructured",":","Whether","to","apply","pruning","globally","on","the","model",".","If","``","parameters_to_prune","``","is","provided",",","global","unstructured","will","be","restricted","on","them",".","amount",":","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",":","Whether","to","apply","pruning",".","-","``","bool","``",".","Always","apply","it","or","not",".","-","``","Callable","[","[","epoch","]",",","bool","]","``",".","For","dynamic","values",".","Will","be","called","every","epoch",".","make_pruning_permanent",":","Whether","to","remove","all","reparametrization","pre-hooks","and","apply","masks","when","training","ends","or","the","model","is","saved",".","use_lottery_ticket_hypothesis",":","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",":","Used","with","``","use_lottery_ticket_hypothesis","``",".","If","True",",","the","model","parameters","will","be","resampled",",","otherwise",",","the","exact","original","parameters","will","be","used",".","pruning_dim",":","If","you","are","using","a","structured","pruning","method","you","need","to","specify","the","dimension",".","pruning_norm",":","If","you","are","using","``","ln_structured","``","you","need","to","specify","the","norm",".","verbose",":","Verbosity","level",".","0","to","disable",",","1","to","log","overall","sparsity",",","2","to","log","per-layer","sparsity","prune_on_train_epoch_end",":","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",":","class",":","`","torch.nn.utils.prune.BasePruningMethod","`",",","or","if","``","amount","``","is","none","of","``","int","``",",","``","float","``","and","``","Callable","``","."]
code: def __init__(
      self,
      pruning_fn: Union[Callable, str],
      parameters_to_prune: _PARAM_LIST = (),
      parameter_names: Optional[List[str]] = None,
      use_global_unstructured: bool = True,
      amount: Union[
          int, float, Callable[[int], Union[int, float]]
      ] = 0.5,
      apply_pruning: Union[bool, Callable[[int], bool]] = True,
      make_pruning_permanent: bool = True,
      use_lottery_ticket_hypothesis: Union[
          bool, Callable[[int], bool]
      ] = True,
      resample_parameters: bool = False,
      pruning_dim: Optional[int] = None,
      pruning_norm: Optional[int] = None,
      verbose: int = 0,
      prune_on_train_epoch_end: bool = True,
  ) -> None:
      """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 :ref:`experimental <versioning:Experimental API>` feature.
  
      .. code-block:: 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.
  
      Args:
  
          pruning_fn: 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: List of tuples ``(nn.Module, "parameter_name_string")``.
  
          parameter_names: List of parameter names to be pruned from the nn.Module.
              Can either be ``"weight"`` or ``"bias"``.
  
          use_global_unstructured: Whether to apply pruning globally on the model.
              If ``parameters_to_prune`` is provided, global unstructured will be restricted on them.
  
          amount: 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: Whether to apply pruning.
  
              - ``bool``. Always apply it or not.
              - ``Callable[[epoch], bool]``. For dynamic values. Will be called every epoch.
  
          make_pruning_permanent: Whether to remove all reparametrization pre-hooks and apply masks
              when training ends or the model is saved.
  
          use_lottery_ticket_hypothesis: 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: Used with ``use_lottery_ticket_hypothesis``. If True, the model parameters will
              be resampled, otherwise, the exact original parameters will be used.
  
          pruning_dim: If you are using a structured pruning method you need to specify the dimension.
  
          pruning_norm: If you are using ``ln_structured`` you need to specify the norm.
  
          verbose: Verbosity level. 0 to disable, 1 to log overall sparsity, 2 to log per-layer sparsity
  
          prune_on_train_epoch_end: 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 :class:`torch.nn.utils.prune.BasePruningMethod`, or
              if ``amount`` is none of ``int``, ``float`` and ``Callable``.
  
      """
  
      self._use_global_unstructured = use_global_unstructured
      self._parameters_to_prune = parameters_to_prune
      self._use_lottery_ticket_hypothesis = (
          use_lottery_ticket_hypothesis
      )
      self._resample_parameters = resample_parameters
      self._prune_on_train_epoch_end = prune_on_train_epoch_end
      self._parameter_names = parameter_names or self.PARAMETER_NAMES
      self._global_kwargs: Dict[str, Any] = {}
      self._original_layers: Optional[Dict[int, _LayerRef]] = None
      self._pruning_method_name: Optional[str] = None
  
      for name in self._parameter_names:
          if name not in self.PARAMETER_NAMES:
              raise MisconfigurationException(
                  f"The provided `parameter_names` name: {name} isn't in {self.PARAMETER_NAMES}"
              )
  
      if isinstance(pruning_fn, str):
          pruning_kwargs = {}
          pruning_fn = pruning_fn.lower()
          if pruning_fn not in _PYTORCH_PRUNING_FUNCTIONS:
              raise MisconfigurationException(
                  f"The provided `pruning_fn` {pruning_fn} isn't available in PyTorch's"
                  f" built-in functions: {list(_PYTORCH_PRUNING_FUNCTIONS.keys())} "
              )
          if pruning_fn.endswith("_structured"):
              if pruning_dim is None:
                  raise MisconfigurationException(
                      "When requesting `structured` pruning, the `pruning_dim` should be provided."
                  )
              if pruning_fn == "ln_structured":
                  if pruning_norm is None:
                      raise MisconfigurationException(
                          "When requesting `ln_structured` pruning, the `pruning_norm` should be provided."
                      )
                  pruning_kwargs["n"] = pruning_norm
              pruning_kwargs["dim"] = pruning_dim
          pruning_fn = self._create_pruning_fn(
              pruning_fn, **pruning_kwargs
          )
      elif self._is_pruning_method(pruning_fn):
          if not use_global_unstructured:
              raise MisconfigurationException(
                  "PyTorch `BasePruningMethod` is currently only supported with `use_global_unstructured=True`."
              )
      else:
          raise MisconfigurationException(
              f"`pruning_fn` is expected to be a str in {list(_PYTORCH_PRUNING_FUNCTIONS.keys())}"
              f" or a PyTorch `BasePruningMethod`. Found: {pruning_fn}."
              " HINT: if passing a `BasePruningMethod`, pass the the class, not an instance"
          )
  
      # need to ignore typing here since pytorch base class does not define the PRUNING_TYPE attribute
      if use_global_unstructured and pruning_fn.PRUNING_TYPE != "unstructured":  # type: ignore
          raise MisconfigurationException(
              'Only the "unstructured" PRUNING_TYPE is supported with `use_global_unstructured=True`.'  # type: ignore
              f" Found method {pruning_fn} of type {pruning_fn.PRUNING_TYPE}. "
          )
  
      self.pruning_fn = pruning_fn
      self._apply_pruning = apply_pruning
      self._make_pruning_permanent = make_pruning_permanent
  
      if not (isinstance(amount, (int, float)) or callable(amount)):
          raise MisconfigurationException(
              "`amount` should be provided and be either an int, a float or Callable function."
          )
  
      self.amount = amount
  
      if verbose not in (0, 1, 2):
          raise MisconfigurationException(
              "`verbose` must be any of (0, 1, 2)"
          )
  
      self._verbose = verbose
  
imports: import inspect
  import logging
  from copy import deepcopy
  from functools import partial
  from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
  import torch.nn.utils.prune
  from lightning_utilities.core.apply_func import apply_to_collection
  from torch import Tensor, nn
  from typing_extensions import TypedDict, override
  import lightning.pytorch
  from lightning.pytorch.callbacks.callback import Callback
  from lightning.pytorch.core.module import LightningModule
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from lightning.pytorch.utilities.rank_zero import rank_zero_debug, rank_zero_only
label_desc: Use image node_id 1 to create a new ModelPruning object from inherited base classes: Callback with example: obj = ModelPruning(pruning_fn, parameters_to_prune, parameter_names, use_global_unstructured, amount, apply_pruning, make_pruning_permanent, use_lottery_ticket_hypothesis, resample_parameters, pruning_dim, pruning_norm, verbose, prune_on_train_epoch_end)