a
    d7Y                     @   sR  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m	Z	m
Z
mZmZmZmZmZ ddlm  m  mZ ddlmZ ddlmZmZ ddl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&e'Z(ej)ej*ej+ej,dZ-ej.ej/ej0ej1dZ2eej3e4f Z5ee5 Z6e ej7ej8ej9fZ:G dd deZ;G dd deZ<dS )z
ModelPruning
^^^^^^^^^^^^
    N)deepcopypartial)AnyCallableDictListOptionalSequenceTupleUnion)apply_to_collection)nnTensor)	TypedDict)Callback)LightningModule)MisconfigurationException)rank_zero_debugrank_zero_only)ln_structuredl1_unstructuredrandom_structuredrandom_unstructuredc                   @   s,   e Zd ZU ejed< eeee	f  ed< dS )	_LayerRefdatanamesN)
__name__
__module____qualname__r   Module__annotations__r   r   intstr r$   r$   l/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/pytorch_lightning/callbacks/pruning.pyr   8   s   

r   c                   @   s|  e Zd ZdZdEeeef eee	e  e
eeeeegeeef f f ee
eege
f f e
ee
eege
f f e
ee ee ee
ddd	d
ZdFeedddZeeeeejf dddZeeeedddZejddddZeejejeddddZddddZeddddZeeeef ddd Zeddd!d"Zeejeeeef d#d$d%Zeeef ddd&d'Z e!dGe	eeef  e	eeef  eeef dd(d)d*Z"d+e#edd,d-d.Z$edd/d0d1Z%d+e#dd2d3d4Z&d+e#dd2d5d6Z'd+e#dd2d7d8Z(e#eeef d9d:d;Z)d+e#eeef dd<d=d>Z*edHe#ee+e ed?d@dAZ,eee
dBdCdDZ-dS )IModelPruning)ZweightZbiasr$   NT      ?Fr   )
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                 C   s  || _ || _|| _|	| _|| _|p&| j| _i | _d| _d| _	| jD ]$}|| jvrBt
d| d| j qBt|tri }| }|tvrt
d| dtt  d|dr|
du rt
d|d	kr|du rt
d
||d< |
|d< | j|fi |}n:| |r|s4t
dnt
dtt  d| d|r^|jdkr^t
d| d|j d|| _|| _|| _t|ttfst|st
d|| _|dvrt
d|| _dS )ao  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:: ``ModelPruning`` is in beta and subject to change.

        .. 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``.
        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:  Z_structuredzKWhen requesting `structured` pruning, the `pruning_dim` should be provided.r   zOWhen requesting `ln_structured` pruning, the `pruning_norm` should be provided.nZ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: zM. HINT: if passing a `BasePruningMethod`, pass the the class, not an instanceZ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_FUNCTIONSlistkeysendswith_create_pruning_fn_is_pruning_methodZPRUNING_TYPEr(   _apply_pruning_make_pruning_permanentr"   floatcallabler,   _verbose)selfr(   r)   r*   r+   r,   r-   r.   r/   r0   r1   r2   r3   r4   nameZpruning_kwargsr$   r$   r%   __init__@   s    `





zModelPruning.__init__)r)   r5   c                 C   s   |S )zAThis function can be overridden to control which module to prune.r$   )rQ   r)   r$   r$   r%   filter_parameters_to_prune   s    z'ModelPruning.filter_parameters_to_prune)r(   kwargsr5   c                 K   sV   | j rt| nt| }t|s&J d| j r2|| _|j| _| j rD|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)	r:   _PYTORCH_PRUNING_METHODrF   rO   rA   r   rC   r&   _wrap_pruning_fn)rQ   r(   rU   Zpruning_methr$   r$   r%   rJ      s    
zModelPruning._create_pruning_fnc                 K   s   t | fi |S Nr   )r(   rU   r$   r$   r%   rW     s    zModelPruning._wrap_pruning_fn)moduler5   c                 C   sN   |  D ]@\}}t|jD ],}|j| }t|tjr|| |j|= qqdS )zRemoves 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)Znamed_modulesrG   Z_forward_pre_hooksrD   pytorch_pruneBasePruningMethodremove)rQ   rY   _khookr$   r$   r%   r.     s    

z#ModelPruning.make_pruning_permanent)newoldrR   r5   c                 C   sP   t | |}t ||}|d u s8|d u s8t|tr8t|ts<d S |j|j|_d S rX   )getattrrD   r   r   toZdevice)r`   ra   rR   dstsrcr$   r$   r%   _copy_param  s
    

$zModelPruning._copy_param)r5   c                 C   s   | j dusJ | j  D ]h}|d }|d }| jrVt|drVt|jrVt|}|  |D ]$\}}| j| \}}| ||| qZqdS )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)	rB   valuesr=   hasattrrO   rg   r   r;   rf   )rQ   dcopyr   irR   r`   new_namer$   r$   r%   apply_lottery_ticket_hypothesis  s    z,ModelPruning.apply_lottery_ticket_hypothesis)r,   r5   c                 C   s$   | j D ]\}}| j|||d qd S )N)rR   r,   )r;   r(   )rQ   r,   rY   rR   r$   r$   r%   _apply_local_pruning1  s    z!ModelPruning._apply_local_pruningc                    s>   || j d< tt| jj  d  fdd| j  D S )Nr,   rQ   c                    s   i | ]\}}| v r||qS r$   r$   ).0r^   vparamsr$   r%   
<dictcomp>9      z7ModelPruning._resolve_global_kwargs.<locals>.<dictcomp>)rA   setinspect	signaturer(   
parametersdiscarditemsrQ   r,   r$   rr   r%   _resolve_global_kwargs5  s    

z#ModelPruning._resolve_global_kwargsc                 C   s$   t j| jfd| ji| | d S )NZpruning_method)rZ   Zglobal_unstructuredr;   r(   r}   r|   r$   r$   r%   _apply_global_pruning;  s    z"ModelPruning._apply_global_pruning)rY   rR   r5   c                 C   s:   | d}t | |sdS t| |}|dk  | fS )N_mask)r   r8   r   )ri   rb   sumitemnumel)rY   rR   attrmaskr$   r$   r%   _get_pruned_stats@  s
    


zModelPruning._get_pruned_statsc                    sd    j r fdd jD } jr, | n
 |  j r` fdd jD } j|||d dS )z+Applies pruning to ``parameters_to_prune``.c                    s   g | ]\}}  ||qS r$   r   rp   mr7   rQ   r$   r%   
<listcomp>K  ru   z.ModelPruning.apply_pruning.<locals>.<listcomp>c                    s   g | ]\}}  ||qS r$   r   r   r   r$   r%   r   S  ru   )r,   N)rP   r;   r:   r~   ro   _log_sparsity_stats)rQ   r,   Z
prev_statsZ
curr_statsr$   r   r%   r-   H  s    
zModelPruning.apply_pruning)prevcurrr,   r5   c                 C   s  t dd | jD }t dd |D }t dd |D }td| j d| d| d|| d	d
| d| d|| d	d | jdkrt| jD ]n\}\}}	|| \}
}|| \}}td| j d|d|	 d| d|
 d|
| d	d
| d|| d	d qd S )Nc                 s   s(   | ] \}}|  D ]}| V  qqd S rX   )ry   r   )rp   Zlayerr]   pr$   r$   r%   	<genexpr>Z  ru   z3ModelPruning._log_sparsity_stats.<locals>.<genexpr>c                 s   s   | ]\}}|V  qd S rX   r$   rp   zerosr]   r$   r$   r%   r   [  ru   c                 s   s   | ]\}}|V  qd S rX   r$   r   r$   r$   r%   r   \  ru   z	Applied `z`. Pruned: /z (z.2%z) -> )r9   z` to `.z` with amount=z
. Pruned: )r   r;   loginforC   rP   	enumerate)rQ   r   r   r,   Ztotal_paramsZprev_total_zerosZcurr_total_zerosrl   rY   rR   Zprev_mask_zerosZprev_mask_sizeZcurr_mask_zerosZcurr_mask_sizer$   r$   r%   r   V  sB    
z ModelPruning._log_sparsity_statsz
pl.Trainer)trainer	pl_modulestager5   c           	      C   s   | j || j| jd}| || _| jr~i | _t| jD ]F\}\}}t|}| j|t	t
|g d | j| d ||f q6d S )N)r*   )r   r   r   )sanitize_parameters_to_pruner;   r@   rT   r<   rB   r   id
setdefaultr   r   append)	rQ   r   r   r   r)   rl   rY   rR   Zid_r$   r$   r%   setupl  s    
zModelPruning.setup)current_epochr5   c                 C   sr   t | jr| |n| j}t | jr.| |n| j}|r<|s@d S | | t | jr`| |rnn| jrn|   d S rX   )rO   rL   r,   r-   r<   rn   )rQ   r   pruner,   r$   r$   r%   _run_pruning|  s    
zModelPruning._run_pruning)r   r   r5   c                 C   s   | j rtd | |j d S )Nz3`ModelPruning.on_train_epoch_end`. Applying pruning)r>   r   r   r   rQ   r   r   r$   r$   r%   on_train_epoch_end  s    zModelPruning.on_train_epoch_endc                 C   s$   |j s | js td | |j d S )Nz8`ModelPruning.on_validation_epoch_end`. Applying pruning)Zsanity_checkingr>   r   r   r   r   r$   r$   r%   on_validation_epoch_end  s    z$ModelPruning.on_validation_epoch_endc                 C   s   | j rtd | | d S )NzJ`ModelPruning.on_train_end`. Pruning is made permanent for this checkpoint)rM   r   r.   r   r$   r$   r%   on_train_end  s    zModelPruning.on_train_end)r   r5   c                 C   sr   |  }dd | D }|D ]6}||d }||d }|j|jd| ||< qttddd}t|t|S )	Nc                 S   s"   h | ]}| d r|d dqS )r    )rI   replace)rp   r^   r$   r$   r%   	<setcomp>  ru   zEModelPruning._make_pruning_permanent_on_state_dict.<locals>.<setcomp>Z_origr   )dtype)tensorr5   c                 S   s   |   S rX   )cpu)r   r$   r$   r%   move_to_cpu  s    zGModelPruning._make_pruning_permanent_on_state_dict.<locals>.move_to_cpu)
state_dictrH   poprc   r   r   r   )rQ   r   r   Zmap_pruned_paramsZtensor_nameorigr   r   r$   r$   r%   %_make_pruning_permanent_on_state_dict  s    z2ModelPruning._make_pruning_permanent_on_state_dict)r   r   
checkpointr5   c                 C   s    | j rtd | ||d< d S )NzP`ModelPruning.on_save_checkpoint`. Pruning is made permanent for this checkpointr   )rM   r   r   )rQ   r   r   r   r$   r$   r%   on_save_checkpoint  s    zModelPruning.on_save_checkpoint)r   r)   r*   r5   c                    s   |pt j}dd |  D  |s4 fdd|D }nt|ttfrt|dkrtdd |D rtdd |D rg g  }}|D ]0\}}| vr|| qt	||s|| q|s|rt
d| d	| nt
d
|S )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.
        c                 S   s   g | ]}t |ts|qS r$   )rD   _MODULE_CONTAINERS)rp   r   r$   r$   r%   r     ru   z=ModelPruning.sanitize_parameters_to_prune.<locals>.<listcomp>c                    s.   g | ]&} D ]}t ||d d ur||fqqS rX   )rb   )rp   r   r   Zcurrent_modulesr$   r%   r     s   r   c                 s   s   | ]}t |d kV  qdS )r9   N)len)rp   r   r$   r$   r%   r     ru   z<ModelPruning.sanitize_parameters_to_prune.<locals>.<genexpr>c                 s   s(   | ] \}}t |tjot |tV  qd S rX   )rD   r   r    r#   )rp   abr$   r$   r%   r     ru   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&   r?   modulesrD   rG   tupler   allr   ri   r   )r   r)   r*   ry   Zmissing_modulesZmissing_parametersrY   rR   r$   r   r%   r     sB    





z)ModelPruning.sanitize_parameters_to_prune)methodr5   c                 C   s   t | sdS t| tjS )NF)rw   isclass
issubclassrZ   r[   )r   r$   r$   r%   rK     s    
zModelPruning._is_pruning_method)r$   NTr'   TTTFNNr   T)r$   )r   )r$   r$   ).r   r   r   r?   r   r   r#   _PARAM_LISTr	   r   boolr"   rN   rS   rT   r   rZ   r[   rJ   staticmethodrW   r   r    r.   rf   rn   ro   r   r}   r~   r   r   r-   r   r   r   r   r   r   r   r   r   r   r
   r   rK   r$   r$   r$   r%   r&   =   s               

 ) & 
.r&   )=__doc__rw   loggingrk   r   	functoolsr   typingr   r   r   r   r	   r
   r   r   Ztorch.nn.utils.pruner   utilsr   rZ   Z#lightning_utilities.core.apply_funcr   Ztorchr   Ztyping_extensionsr   Zpytorch_lightningplZ$pytorch_lightning.callbacks.callbackr   Zpytorch_lightning.core.moduler   Z&pytorch_lightning.utilities.exceptionsr   Z%pytorch_lightning.utilities.rank_zeror   r   	getLoggerr   r   r   r   r   r   rF   ZLnStructuredZL1UnstructuredZRandomStructuredZRandomUnstructuredrV   r    r#   Z_PARAM_TUPLEr   Z
SequentialZ
ModuleListZ
ModuleDictr   r   r&   r$   r$   r$   r%   <module>   s<   (
