Ë
    ÿÍ:jhZ  ã                   óT  — d Z ddlZddlZddlmZ ddlmZ ddlmZ ddl	m
Z
mZmZmZ ddlmc mc mZ ddlmZ ddlmZmZ dd	lmZmZ ddlZdd
lmZ ddlmZ ddl m!Z! ddl"m#Z#m$Z$  ejJ                  e&«      Z'ejP                  ejR                  ejT                  ejV                  dœZ,ejZ                  ej\                  ej^                  ej`                  dœZ1e2ejf                  e4f   Z5ee5   Z6eejn                  ejp                  ejr                  fZ: G d„ de«      Z; G d„ de«      Z<y)z
ModelPruning
^^^^^^^^^^^^
é    N)ÚSequence)Údeepcopy©Úpartial)ÚAnyÚCallableÚOptionalÚUnion)Úapply_to_collection)ÚTensorÚnn)Ú	TypedDictÚoverride)ÚCallback)ÚLightningModule)ÚMisconfigurationException)Úrank_zero_debugÚrank_zero_only)Úln_structuredÚl1_unstructuredÚrandom_structuredÚrandom_unstructuredc                   óF   — e Zd ZU ej                  ed<   eeee	f      ed<   y)Ú	_LayerRefÚdataÚnamesN)
Ú__name__Ú
__module__Ú__qualname__r   ÚModuleÚ__annotations__ÚlistÚtupleÚintÚstr© ó    úx/home/mcse/projects/srt_converter/srt-converter-venv/lib/python3.12/site-packages/pytorch_lightning/callbacks/pruning.pyr   r   :   s   … Ø
�)‰)ƒOØ��c˜3�h‘Ñ Ô r'   r   c                   óª  — e Zd ZdZ	 	 	 	 	 	 	 	 	 	 	 	 d5deeef   dedee	e      de
deeeeegeeef   f   f   dee
eege
f   f   d	e
d
ee
eege
f   f   de
dee   dee   dede
ddfd„Zd6dedefd„Zdededeeej"                  f   fd„Zedededefd„«       Zdej,                  ddfd„Zedej,                  dej,                  deddfd„«       Zd7d„Zdeddfd„Zdedeeef   fd„Zdeddfd„Zedej,                  dedeeef   fd „«       Zdeeef   ddfd!„Z e!	 d8d"e	eeef      d#e	eeef      deeef   ddfd$„«       Z"e#d%d&d'e$d(eddfd)„«       Z%d*eddfd+„Z&e#d%d&d'e$ddfd,„«       Z'e#d%d&d'e$ddfd-„«       Z(e#d%d&d'e$ddfd.„«       Z)d'e$deeef   fd/„Z*e#d%d&d'e$d0eeef   ddfd1„«       Z+e	 d9d'e$dede,e   defd2„«       Z-ed3ede
fd4„«       Z.y):ÚModelPruning)ÚweightÚbiasNÚ
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Úreturnc           	      ó  — || _         || _        || _        |	| _        || _        |xs | j
                  | _        i | _        d| _        d| _	        | j                  D ]+  }|| j
                  vsŒt        d|› d| j
                  › �«      ‚ t        |t        «      r•i }|j                  «       }|t        vr-t        d|› dt        t        j!                  «       «      › d�«      ‚|j#                  d«      r)|
€t        d«      ‚|d	k(  r|€t        d
«      ‚||d<   |
|d<    | j$                  |fi |¤Ž}nK| j'                  |«      r|s8t        d«      ‚t        dt        t        j!                  «       «      › d|› d�«      ‚|r+|j(                  dk7  rt        d|› d|j(                  › d�«      ‚|| _        || _        || _        t        |t0        t2        f«      st5        |«      st        d«      ‚|| _        |dvrt        d«      ‚|| _        y)a…  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 reparameterization 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``.

        Nz%The provided `parameter_names` name: z
 isn't in zThe provided `pruning_fn` z2 isn't available in PyTorch's built-in functions: ú Ú_structuredzKWhen requesting `structured` pruning, the `pruning_dim` should be provided.r   zOWhen requesting `ln_structured` pruning, the `pruning_norm` should be provided.ÚnÚdimz\PyTorch `BasePruningMethod` is currently only supported with `use_global_unstructured=True`.z(`pruning_fn` is expected to be a str in z* or a PyTorch `BasePruningMethod`. Found: zI. HINT: if passing a `BasePruningMethod`, pass the class, not an instanceÚunstructuredzdOnly the "unstructured" PRUNING_TYPE is supported with `use_global_unstructured=True`. Found method z	 of type z. zO`amount` should be provided and be either an int, a float or Callable function.)r   é   é   z"`verbose` must be any of (0, 1, 2))Ú_use_global_unstructuredÚ_parameters_to_pruneÚ_use_lottery_ticket_hypothesisÚ_resample_parametersÚ_prune_on_train_epoch_endÚPARAMETER_NAMESÚ_parameter_namesÚ_global_kwargsÚ_original_layersÚ_pruning_method_namer   Ú
isinstancer%   ÚlowerÚ_PYTORCH_PRUNING_FUNCTIONSr"   ÚkeysÚendswithÚ_create_pruning_fnÚ_is_pruning_methodÚPRUNING_TYPEr-   Ú_apply_pruningÚ_make_pruning_permanentr$   ÚfloatÚcallabler1   Ú_verbose)Úselfr-   r.   r/   r0   r1   r2   r3   r4   r5   r6   r7   r8   r9   ÚnameÚpruning_kwargss                   r(   Ú__init__zModelPruning.__init__B   sm  € ðB )@ˆÔ%Ø$7ˆÔ!Ø.KˆÔ+Ø$7ˆÔ!Ø)AˆÔ&Ø /Ò G°4×3GÑ3GˆÔØ.0ˆÔØ@DˆÔØ37ˆÔ!à×)Ñ)ò 	ˆDØ˜4×/Ñ/Ò/Ü/Ø;¸D¸6ÀÈD×L`ÑL`ÐKaÐbóð ð	ô �j¤#Ô&ØˆNØ#×)Ñ)Ó+ˆJØÔ!;Ñ;Ü/Ø0°°ð =,Ü,0Ô1K×1PÑ1PÓ1RÓ,SÐ+TÐTUðWóð ð ×"Ñ" =Ô1ØÐ&Ü3Øeóð ð  Ò0Ø#Ð+Ü7Ømóð ð +7�N 3Ñ'Ø(3�˜uÑ%Ø0˜×0Ñ0°ÑN¸~ÑN‰JØ×$Ñ$ ZÔ0Ù*Ü/Øróð ô ,Ø:¼4Ô@Z×@_Ñ@_Ó@aÓ;bÐ:cØ<¸Z¸Lð I[ð[óð ñ # z×'>Ñ'>À.Ò'PÜ+ð!Ø!+ ¨I°j×6MÑ6MÐ5NÈbðRóð ð
 %ˆŒØ+ˆÔØ'=ˆÔ$ä˜6¤C¬ <Ô0´H¸VÔ4DÜ+Øaóð ð ˆŒà˜)Ñ#Ü+Ð,PÓQÐQàˆ�r'   c                 ó   — |S )zAThis function can be overridden to control which module to prune.r&   )rZ   r.   s     r(   Úfilter_parameters_to_prunez'ModelPruning.filter_parameters_to_pruneë   s   € à"Ð"r'   Úkwargsc                 óò   — | j                   r	t        |   nt        |   }t        |«      sJ d«       ‚| j                   r|| _        |j
                  | _        | j                   r|S t        j                  |fi |¤ŽS )a  This function takes `pruning_fn`, a function name.

        IF use_global_unstructured, pruning_fn will be resolved into its associated ``PyTorch BasePruningMethod`` ELSE,
        pruning_fn will be resolved into its function counterpart from `torch.nn.utils.prune`.

        z'Selected pruning method is not callable)	rC   Ú_PYTORCH_PRUNING_METHODrO   rX   rJ   r   rL   r*   Ú_wrap_pruning_fn)rZ   r-   r`   Úpruning_meths       r(   rR   zModelPruning._create_pruning_fnï   s€   € ð ×,Ò,ô $ JÒ/ä+¨JÑ7ð 	ô
 ˜Ô%ÐPÐ'PÓPÐ%Ø×(Ò(Ø"(ˆDÔð %1×$9Ñ$9ˆÔ!Ø×(Ò(ØÐÜ×,Ñ,¨\ÑD¸VÑDÐDr'   c                 ó   — t        | fi |¤ŽS ©Nr   )r-   r`   s     r(   rc   zModelPruning._wrap_pruning_fn  s   € ä�zÑ, VÑ,Ð,r'   Úmodulec                 óø   — |j                  «       D ]g  \  }}t        |j                  «      D ]J  }|j                  |   }t        |t        j
                  «      sŒ-|j                  |«       |j                  |= ŒL Œi y)z£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

        N)Únamed_modulesr"   Ú_forward_pre_hooksrM   Úpytorch_pruneÚBasePruningMethodÚremove)rZ   rg   Ú_ÚkÚhooks        r(   r3   z#ModelPruning.make_pruning_permanent	  st   € ð  ×-Ñ-Ó/ò 	5‰IˆAˆvÜ˜&×3Ñ3Ó4ò 5�Ø×0Ñ0°Ñ3�Ü˜d¤M×$CÑ$CÕDØ—K‘K Ô'Ø×1Ñ1°!Ñ4ñ	5ñ	5r'   ÚnewÚoldr[   c                 ó  — t        | |dz   «      rt        | |dz   «      nt        | |«      }t        ||«      }|�"|� t        |t        «      rt        |t        «      sy |j                  j                  |j                  «      |_        y )NÚ_orig)ÚhasattrÚgetattrrM   r   r   ÚtoÚdevice)rq   rr   r[   ÚdstÚsrcs        r(   Ú_copy_paramzModelPruning._copy_param  sq   € ô /6°c¸4À'¹>Ô.JŒg�c˜4 '™>Ô*ÔPWÐX[Ð]aÓPbˆÜ�c˜4Ó ˆØˆ;˜#˜+¬Z¸¼VÔ-DÌJÐWZÔ\bÔLcØØ—8‘8—;‘;˜sŸz™zÓ*ˆ�r'   c                 ó`  — | j                   €J ‚| j                   j                  «       D ]ƒ  }|d   }|d   }| j                  r<t        |d«      r0t	        |j
                  «      rt        |«      }|j                  «        |D ]*  \  }}| j                  |   \  }}| j                  |||«       Œ, Œ… y)aÐ  Lottery ticket hypothesis algorithm (see page 2 of the paper):

            1. Randomly initialize a neural network :math:`f(x; \theta_0)` (where :math:`\theta_0 \sim \mathcal{D}_\theta`).
            2. Train the network for :math:`j` iterations, arriving at parameters :math:`\theta_j`.
            3. Prune :math:`p\%` of the parameters in :math:`\theta_j`, creating a mask :math:`m`.
            4. Reset the remaining parameters to their values in :math:`\theta_0`, creating the winning ticket :math:`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 :math:`\theta_z \sim \mathcal{D}_\theta`

        Nr   r   Úreset_parameters)	rK   ÚvaluesrF   ru   rX   r}   r   rD   r{   )rZ   ÚdÚcopyr   Úir[   rq   rn   s           r(   Úapply_lottery_ticket_hypothesisz,ModelPruning.apply_lottery_ticket_hypothesis  s¶   € ð ×$Ñ$Ð0Ð0Ð0Ø×&Ñ&×-Ñ-Ó/ò 	2ˆAØ�V‘9ˆDØ�g‘JˆEØ×(Ò(¬W°TÐ;MÔ-NÔS[Ð\`×\qÑ\qÔSrÜ “~�Ø×%Ñ%Ô'Ø ò 2‘��4Ø×2Ñ2°1Ñ5‘��QØ× Ñ   d¨DÕ1ñ2ñ	2r'   c                 óT   — | j                   D ]  \  }}| j                  |||¬«       Œ y )N)r[   r1   )rD   r-   )rZ   r1   rg   r[   s       r(   Ú_apply_local_pruningz!ModelPruning._apply_local_pruning7  s-   € Ø ×5Ñ5ò 	>‰LˆF�DØ�O‰O˜F¨°fˆOÕ=ñ	>r'   c                 ó  — || j                   d<   t        t        j                  | j                  «      j
                  «      }|j                  d«       | j                   j                  «       D ��ci c]  \  }}||v sŒ||“Œ c}}S c c}}w )Nr1   rZ   )rJ   ÚsetÚinspectÚ	signaturer-   Ú
parametersÚdiscardÚitems)rZ   r1   Úparamsro   Úvs        r(   Ú_resolve_global_kwargsz#ModelPruning._resolve_global_kwargs;  sn   € Ø(.ˆ×Ñ˜HÑ%Ü”W×&Ñ& t§¡Ó7×BÑBÓCˆØ�‰�vÔØ!%×!4Ñ!4×!:Ñ!:Ó!<×L™˜˜AÀÀVÂ��1‘ÓLÐLùÓLs   Á0BÁ=Bc                 óz   — t        j                  | j                  fd| j                  i| j	                  |«      ¤Ž y )NÚpruning_method)rk   Úglobal_unstructuredrD   r-   rŽ   )rZ   r1   s     r(   Ú_apply_global_pruningz"ModelPruning._apply_global_pruningA  s9   € Ü×)Ñ)Ø×%Ñ%ñ	
Ø6:·o±oð	
ØIM×IdÑIdÐekÓIló	
r'   c                 ó    — |› d�}t        | |«      syt        | |«      }|dk(  j                  «       j                  «       |j	                  «       fS )NÚ_mask)r   rA   r   )ru   rv   ÚsumÚitemÚnumel)rg   r[   ÚattrÚmasks       r(   Ú_get_pruned_statszModelPruning._get_pruned_statsF  sK   € à��uˆ~ˆÜ�v˜tÔ$ØÜ�v˜tÓ$ˆØ˜‘	�‰Ó ×%Ñ%Ó'¨¯©«Ð5Ð5r'   c                 óŒ  — | j                   r.| j                  D ��cg c]  \  }}| j                  ||«      ‘Œ }}}| j                  r| j	                  |«       n| j                  |«       | j                   rC| j                  D ��cg c]  \  }}| j                  ||«      ‘Œ }}}| j                  ||¬«       yyc c}}w c c}}w )z+Applies pruning to ``parameters_to_prune``.)r1   N)rY   rD   rš   rC   r’   r„   Ú_log_sparsity_stats)rZ   r1   Úmr>   Ú
prev_statsÚ
curr_statss         r(   r2   zModelPruning.apply_pruningN  s¯   € à�=Š=ØCG×C\ÑC\×]¹4¸1¸a˜$×0Ñ0°°AÕ6Ð]ˆJÑ]à×(Ò(Ø×&Ñ& vÕ.à×%Ñ% fÔ-à�=Š=ØCG×C\ÑC\×]¹4¸1¸a˜$×0Ñ0°°AÕ6Ð]ˆJÑ]Ø×$Ñ$ Z°ÀFÐ$ÕKð ùó ^ùó ^s   œB:ÂC ÚprevÚcurrc                 óê  — t        d„ |D «       «      }t        d„ |D «       «      }t        d„ |D «       «      }t        j                  d| j                  › d|› d|› d||z  d›d	|› d|› d||z  d›d
�«       | j                  dk(  rqt        | j                  «      D ]X  \  }\  }}	||   \  }
}||   \  }}t        j                  d| j                  › d|›d|	› d|› d|
› d|
|z  d›d	|› d||z  d›d
�«       ŒZ y y )Nc              3   ó&   K  — | ]	  \  }}|–— Œ y ­wrf   r&   )Ú.0rn   Útotals      r(   ú	<genexpr>z3ModelPruning._log_sparsity_stats.<locals>.<genexpr>`  s   è ø€ Ò6¡X Q¨œ5Ñ6ùó   ‚c              3   ó&   K  — | ]	  \  }}|–— Œ y ­wrf   r&   ©r¤   Úzerosrn   s      r(   r¦   z3ModelPruning._log_sparsity_stats.<locals>.<genexpr>a  ó   è ø€ Ò:©¨°œuÑ:ùr§   c              3   ó&   K  — | ]	  \  }}|–— Œ y ­wrf   r&   r©   s      r(   r¦   z3ModelPruning._log_sparsity_stats.<locals>.<genexpr>b  r«   r§   z	Applied `z`. Pruned: ú/z (z.2%z) -> ú)rB   z` to `ú.z` with amount=z
. Pruned: )r•   ÚlogÚinforL   rY   Ú	enumeraterD   )rZ   r    r¡   r1   Útotal_paramsÚprev_total_zerosÚcurr_total_zerosr�   rg   r[   Úprev_mask_zerosÚprev_mask_sizeÚcurr_mask_zerosÚcurr_mask_sizes                 r(   rœ   z ModelPruning._log_sparsity_stats\  sZ  € ô Ñ6°Ô6Ó6ˆÜÑ:°TÔ:Ó:ÐÜÑ:°TÔ:Ó:ÐÜ�‰Ø˜×1Ñ1Ð2ð 3Ø Ð!  < .°Ð3CÀlÑ3RÐSVÐ2Wð XØ Ð!  < .°Ð3CÀlÑ3RÐSVÐ2WÐWXðZô	
ð
 �=‰=˜AÒÜ%.¨t×/HÑ/HÓ%Iò Ñ!�‘>�F˜DØ26°q±'Ñ/� Ø26°q±'Ñ/� Ü—‘Ø × 9Ñ 9Ð:¸&ÀÀ
È!ÈDÈ6ÐQ_Ð`fÐ_gð hØ'Ð(¨¨?¸^Ñ+KÈCÐ*Pð QØ'Ð(¨¨?¸^Ñ+KÈCÐ*PÐPQðSõñð r'   Útrainerz
pl.TrainerÚ	pl_moduleÚstagec           	      ó¤  — | j                  || j                  | j                  ¬«      }| j                  |«      | _        | j                  r†i | _        t        | j                  «      D ]f  \  }\  }}t        |«      }| j
                  j                  |t        t        |«      g ¬«      «       | j
                  |   d   j                  ||f«       Œh y y )N)r/   )r   r   r   )Úsanitize_parameters_to_prunerD   rI   r_   rE   rK   r²   ÚidÚ
setdefaultr   r   Úappend)	rZ   rº   r»   r¼   r.   r�   rg   r[   Úid_s	            r(   ÚsetupzModelPruning.setupr  sÍ   € à"×?Ñ?Ø�t×0Ñ0À$×BWÑBWð @ó 
Ðð %)×$CÑ$CÐDWÓ$XˆÔ!à×.Ò.ð %'ˆDÔ!Ü%.¨t×/HÑ/HÓ%Iò FÑ!�‘>�F˜DÜ˜“j�Ø×%Ñ%×0Ñ0°´iÄXÈfÓEUÐ]_Ô6`ÔaØ×%Ñ% cÑ*¨7Ñ3×:Ñ:¸A¸t¸9ÕEñFð	 /r'   Úcurrent_epochc                 ó‚  — t        | j                  «      r| j                  |«      n| j                  }t        | j                  «      r| j                  |«      n| j                  }|r|sy | j                  |«       t        | j                  «      r| j	                  |«      rn| j                  r| j                  «        y y y rf   )rX   rU   r1   r2   rE   r‚   )rZ   rÄ   Úpruner1   s       r(   Ú_run_pruningzModelPruning._run_pruningƒ  sœ   € Ü6>¸t×?RÑ?RÔ6S�×#Ñ# MÔ2ÐY]×YlÑYlˆÜ/7¸¿¹Ô/D�—‘˜]Ô+È$Ï+É+ˆÙ™FØØ×Ñ˜6Ô"ô ˜×;Ñ;Ô<ð ×/Ñ/°Õ>à×4Ò4à×0Ñ0Õ2ð 5ð ?r'   c                 ój   — | j                   r't        d«       | j                  |j                  «       y y )Nz3`ModelPruning.on_train_epoch_end`. Applying pruning)rG   r   rÇ   rÄ   ©rZ   rº   r»   s      r(   Úon_train_epoch_endzModelPruning.on_train_epoch_end‘  s-   € à×)Ò)ÜÐQÔRØ×Ñ˜i×5Ñ5Õ6ð *r'   c                 ó„   — |j                   s4| j                  s't        d«       | j                  |j                  «       y y y )Nz8`ModelPruning.on_validation_epoch_end`. Applying pruning)Úsanity_checkingrG   r   rÇ   rÄ   rÉ   s      r(   Úon_validation_epoch_endz$ModelPruning.on_validation_epoch_end—  s9   € à×&Ò&¨t×/MÒ/MÜÐVÔWØ×Ñ˜i×5Ñ5Õ6ð 0NÐ&r'   c                 óV   — | j                   rt        d«       | j                  |«       y y )NzJ`ModelPruning.on_train_end`. Pruning is made permanent for this checkpoint)rV   r   r3   rÉ   s      r(   Úon_train_endzModelPruning.on_train_end�  s'   € à×'Ò'ÜÐhÔiØ×'Ñ'¨	Õ2ð (r'   c                 ór  — |j                  «       }|D �ch c]&  }|j                  d«      sŒ|j                  dd«      ’Œ( }}|D ]L  }|j                  |dz   «      }|j                  |dz   «      }|j	                  |j
                  ¬«      |z  ||<   ŒN dt        dt        fd„}t        |t        |«      S c c}w )Nr”   Ú rt   )ÚdtypeÚtensorr:   c                 ó"   — | j                  «       S rf   )Úcpu)rÓ   s    r(   Úmove_to_cpuzGModelPruning._make_pruning_permanent_on_state_dict.<locals>.move_to_cpu®  s   € à—:‘:“<Ðr'   )Ú
state_dictrQ   ÚreplaceÚpoprw   rÒ   r   r   )	rZ   r»   r×   ro   Úmap_pruned_paramsÚtensor_nameÚorigr™   rÖ   s	            r(   Ú%_make_pruning_permanent_on_state_dictz2ModelPruning._make_pruning_permanent_on_state_dict£  s½   € Ø×)Ñ)Ó+ˆ
ð >HÖ_¸È1Ï:É:ÐV]ÕK^˜QŸY™Y w°Õ3Ð_ÐÐ_Ø,ò 	GˆKØ—>‘> +°Ñ"7Ó8ˆDØ—>‘> +°Ñ"7Ó8ˆDà&*§g¡g°D·J±J gÓ&?À$Ñ&FˆJ�{Ò#ð		Gð	 ¤ð 	 ¬6ó 	 ô # :¬v°{ÓCÐCùò `s
   •B4¬B4Ú
checkpointc                 ó\   — | j                   r t        d«       | j                  |«      |d<   y y )NzP`ModelPruning.on_save_checkpoint`. Pruning is made permanent for this checkpointr×   )rV   r   rÝ   )rZ   rº   r»   rÞ   s       r(   Úon_save_checkpointzModelPruning.on_save_checkpoint´  s/   € à×'Ò'ÜÐnÔoà'+×'QÑ'QÐR[Ó'\ˆJ�|Ò$ð (r'   c                 ó”  — |xs t         j                  }| j                  «       D �cg c]  }t        |t        «      rŒ|‘Œ }}|sN|D ��cg c]?  }|D ]8  }t        ||d«      �)t        t        ||d«      t        j                  «      r||f‘Œ: ŒA }}}|S t        |t        t        f«      r‹t        |«      dkD  r}t        d„ |D «       «      rkt        d„ |D «       «      rYg g }}|D ]9  \  }	}
|	|vr|j                  |	«       Œt        |	|
«      rŒ)|j                  |
«       Œ; |s|rt        d|› d|› �«      ‚|S t        d«      ‚c c}w c c}}w )a”  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.

        Nr   c              3   ó8   K  — | ]  }t        |«      d k(  –— Œ y­w)rB   N)Úlen)r¤   Úps     r(   r¦   z<ModelPruning.sanitize_parameters_to_prune.<locals>.<genexpr>Ö  s   è ø€ Ò= A”C˜“F˜a•KÑ=ùs   ‚c              3   óz   K  — | ]3  \  }}t        |t        j                  «      xr t        |t        «      –— Œ5 y ­wrf   )rM   r   r    r%   )r¤   ÚaÚbs      r(   r¦   z<ModelPruning.sanitize_parameters_to_prune.<locals>.<genexpr>×  s.   è ø€ ÒdÉÈÈ1”J˜q¤"§)¡)Ó,ÒC´¸A¼sÓ1CÓCÑdùs   ‚9;zUSome provided `parameters_to_prune` don't exist in the model. Found missing modules: z and missing parameters: zThe provided `parameters_to_prune` should either be list of tuple with 2 elements: (nn.Module, parameter_name_to_prune) or None)r*   rH   ÚmodulesrM   Ú_MODULE_CONTAINERSrv   r   Ú	Parameterr"   r#   rã   ÚallrÁ   ru   r   )r»   r.   r/   r‰   r�   Úcurrent_modulesrä   Úmissing_modulesÚmissing_parametersrg   r[   s              r(   r¾   z)ModelPruning.sanitize_parameters_to_prune»  s‰  € ð %ÒD¬×(DÑ(Dˆ
à&/×&7Ñ&7Ó&9Öc ÄÈAÔOaÕAbš1ÐcˆÐcá"ð $÷#àØ(ò#ð Ü˜1˜a Ó&Ð2´zÄ'È!ÈQÐPTÓBUÔWY×WcÑWcÔ7dð �A’ð#Øð#Ðñ #ð> #Ð"ô1 Ð*¬T´5¨MÔ:ÜÐ'Ó(¨1Ò,ÜÑ=Ð)<Ô=Ô=ÜÑdÐPcÔdÔdà24°bÐ/ˆOØ 3ò 4‘�˜Ø Ñ0Ø#×*Ñ*¨6Ô2ØÜ˜v tÕ,Ø&×-Ñ-¨dÕ3ð4ñ Ñ"4Ü/ð/Ø/>Ð.?Ð?XÐYkÐXlðnóð ð #Ð"ô ,ðQóð ùò; dùó#s   §D?½D?Á
AEÚmethodc                 ób   — t        j                  | «      syt        | t        j                  «      S )NF)r‡   ÚisclassÚ
issubclassrk   rl   )rï   s    r(   rS   zModelPruning._is_pruning_methodî  s$   € ä�‰˜vÔ&ØÜ˜&¤-×"AÑ"AÓBÐBr'   )r&   NTg      à?TTTFNNr   T)r&   )r:   N)r   )r&   r&   )/r   r   r   rH   r
   r   r%   Ú_PARAM_LISTr	   r"   Úboolr$   rW   r]   r_   r   rk   rl   rR   Ústaticmethodrc   r   r    r3   r{   r‚   r„   ÚdictrŽ   r’   r#   rš   r2   r   rœ   r   r   rÃ   rÇ   rÊ   rÍ   rÏ   rÝ   rà   r   r¾   rS   r&   r'   r(   r*   r*   ?   sy  „ Ø(€Oð
 ,.Ø/3Ø(,ØHKØ<@Ø'+ØLPØ$)Ø%)Ø&*ØØ)-ñg à˜( C˜-Ñ(ðg ð )ðg ð " $ s¡)Ñ,ð	g ð
 "&ðg ð �c˜5 (¨C¨5°%¸¸U¸
Ñ2CÐ+CÑ"DÐDÑEðg ð ˜T 8¨S¨E°4¨KÑ#8Ð8Ñ9ðg ð !%ðg ð (-¨T°8¸S¸EÀ4¸KÑ3HÐ-HÑ'Iðg ð "ðg ð ˜c‘]ðg ð ˜s‘mðg ð ðg ð #'ðg ð 
óg ñR#¸kð #ÐS^ó #ðE¨Sð E¸Cð EÀEÈ(ÐTa×TsÑTsÐJsÑDtó Eð, ð- Xð -¸ð -Àò -ó ð-ð5¨R¯Y©Yð 5¸4ó 5ð ð+˜Ÿ™ð +¨¯©ð +¸#ð +À$ò +ó ð+ó2ð0>¨5ð >°Tó >ðM¨Uð M°t¸CÀ¸H±~ó Mð
¨Eð 
°dó 
ð
 ð6 "§)¡)ð 6°3ð 6¸5ÀÀcÀ¹?ò 6ó ð6ðL E¨#¨u¨*Ñ$5ð L¸$ó Lð àdeñØ˜˜s C˜x™Ñ)ðØ15°e¸CÀ¸H±oÑ1FðØPUÐVYÐ[`ÐV`ÑPaðà	òó ðð* ðF˜\ð F°oð FÈcð FÐVZò Fó ðFð 3¨#ð 3°$ó 3ð ð7¨,ð 7À?ð 7ÐW[ò 7ó ð7ð
 ð7¨|ð 7Èð 7Ð\`ò 7ó ð7ð
 ð3 Lð 3¸_ð 3ÐQUò 3ó ð3ð
D¸ð DÐSWÐX[Ð]`ÐX`ÑSaó Dð" ð]¨,ð ]À?ð ]Ð`dÐehÐjmÐemÑ`nð ]Ðswò ]ó ð]ð àlnñ0#Ø"ð0#Ø9Dð0#Ø\dÐehÑ\ið0#à	ò0#ó ð0#ðd ðC 3ð C¨4ò Có ñCr'   r*   )=Ú__doc__r‡   ÚloggingÚcollections.abcr   r€   r   Ú	functoolsr   Útypingr   r   r	   r
   Útorch.nn.utils.pruner   ÚutilsrÆ   rk   Ú#lightning_utilities.core.apply_funcr   Útorchr   Útyping_extensionsr   r   Úpytorch_lightningÚplÚ$pytorch_lightning.callbacks.callbackr   Úpytorch_lightning.core.moduler   Ú&pytorch_lightning.utilities.exceptionsr   Ú%pytorch_lightning.utilities.rank_zeror   r   Ú	getLoggerr   r°   r   r   r   r   rO   ÚLnStructuredÚL1UnstructuredÚRandomStructuredÚRandomUnstructuredrb   r#   r    r%   Ú_PARAM_TUPLEró   Ú
SequentialÚ
ModuleListÚ
ModuleDictré   r   r*   r&   r'   r(   ú<module>r     s  ðñó
 Û Ý $Ý Ý ß 1Ó 1ç ,Ó ,Ý Cß ß 1ã Ý 9Ý 9Ý Lß Qà€g×Ñ˜Ó!€ð #×0Ñ0Ø$×4Ñ4Ø&×8Ñ8Ø(×<Ñ<ñ	Ð ð #×/Ñ/Ø$×3Ñ3Ø&×7Ñ7Ø(×;Ñ;ñ	Ð ð �R—Y‘Y �^Ñ$€Ø�|Ñ$€Ø% r§}¡}°b·m±mÀRÇ]Á]ÐSÐ ô!�	ô !ô
sC�8õ sCr'   