
      iTR                        d Z ddlZddlmZmZ ddlmZmZmZm	Z	 ddl
Z
ddlmZmZ ddlmZ ddl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 esdgZ ej>                  e       Z!de"de#fdZ$ G d de      Z% G d de%      Z&y)z_
Finetuning Callback
^^^^^^^^^^^^^^^^^^^^

Freeze and unfreeze models for finetuning purposes.
    N)	GeneratorIterable)AnyCallableOptionalUnion)Module
ModuleDict)
_BatchNorm)	Optimizer)override)Callback)MisconfigurationException)_TORCHVISION_AVAILABLE)rank_zero_warnBackboneFinetuningepochreturnc                      y)Ng       @ )r   s    {/Volumes/fast/ai/experiments/voice-extract-mac/.venv/lib/python3.12/site-packages/pytorch_lightning/callbacks/finetuning.pymultiplicativer   ,   s        c                      e Zd ZdZd)dZedeeef   fd       Z	edeeef   ddfd       Z
ed*d       Zedeeeeeef      f   dee   fd       Ze	 d+deeeeeef      f   dededefd       Zedeeeeeef      f   ddfd       Zededdfd       Zed,deeeeeef      f   deddfd       Zedededefd       Ze	 	 	 d-deeeeeef      f   dedee   dededdfd       Zedd	d
ddeddfd       Zedeeeef      dedeeeef      fd        Zd
dd!ed"ed#eeeef      ddf
d$Z ed*d%       Z!d
dd&ededdfd'Z"d.d(Z#y)/BaseFinetuninga  This class implements the base logic for writing your own Finetuning Callback.

    .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.

    Override ``freeze_before_training`` and ``finetune_function`` methods with your own logic.

    ``freeze_before_training``: This method is called before ``configure_optimizers``
        and should be used to freeze any modules parameters.

    ``finetune_function``: This method is called on every train epoch start and should be used to
        ``unfreeze`` any parameters. Those parameters need to be added in a new ``param_group``
        within the optimizer.

    .. note:: Make sure to filter the parameters based on ``requires_grad``.

    Example::

        >>> from torch.optim import Adam
        >>> class MyModel(pl.LightningModule):
        ...     def configure_optimizer(self):
        ...         # Make sure to filter the parameters based on `requires_grad`
        ...         return Adam(filter(lambda p: p.requires_grad, self.parameters()))
        ...
        >>> class FeatureExtractorFreezeUnfreeze(BaseFinetuning):
        ...     def __init__(self, unfreeze_at_epoch=10):
        ...         super().__init__()
        ...         self._unfreeze_at_epoch = unfreeze_at_epoch
        ...
        ...     def freeze_before_training(self, pl_module):
        ...         # freeze any module you want
        ...         # Here, we are freezing `feature_extractor`
        ...         self.freeze(pl_module.feature_extractor)
        ...
        ...     def finetune_function(self, pl_module, current_epoch, optimizer):
        ...         # When `current_epoch` is 10, feature_extractor will start training.
        ...         if current_epoch == self._unfreeze_at_epoch:
        ...             self.unfreeze_and_add_param_group(
        ...                 modules=pl_module.feature_extractor,
        ...                 optimizer=optimizer,
        ...                 train_bn=True,
        ...             )

    r   Nc                      i | _         d| _        y NF)_internal_optimizer_metadata_restartingselfs    r   __init__zBaseFinetuning.__init__]   s    MO) r   c                     d| j                   iS )Ninternal_optimizer_metadata)r   r    s    r   
state_dictzBaseFinetuning.state_dicta   s     *4+L+L
 	
r   r%   c                 >    d| _         d|v r|d   | _        y || _        y )NTr$   )r   r   )r!   r%   s     r   load_state_dictzBaseFinetuning.load_state_dictg   s*    (J60:;X0YD- 1;D-r   trainer
pl.Trainer	pl_modulepl.LightningModulec                    | j                   r| j                  rkt        |j                               }t	        |j
                        D ]:  \  }}|| j                  v s| j                  | j                  |   |      }||_        < d| _         y y r   )r   r   dictnamed_parameters	enumerate
optimizers_apply_mapping_to_param_groupsparam_groups)r!   r(   r*   r.   opt_idx	optimizerr2   s          r   on_fit_startzBaseFinetuning.on_fit_startp   s     00#'	(B(B(D#E *3G4F4F*G&GY$"C"CC'+'J'J ==gFHX( 2>	. +H  %D r   modulesc                 j   t        | t              r| j                         } t        | t              r9g }| D ]&  }|j	                  t
        j                  |             ( t        |      }n| j                         }|D cg c]*  }t        |j                               r|j                  s)|, c}S c c}w )aH  This function is used to flatten a module or an iterable of modules into a list of its leaf modules (modules
        with no children) and parent modules that have parameters directly themselves.

        Args:
            modules: A given module or an iterable of modules

        Returns:
            List of modules

        )
isinstancer
   valuesr   extendr   flatten_modulesiterr6   listchildren_parameters)r6   _flatten_modulesm_moduless       r   r;   zBaseFinetuning.flatten_modules~   s     gz*nn&Ggx(! ''(F(Fq(IJ  ,-H(H $O8a4

+=8OOOs   >*B0)B0train_bnrequires_gradc              #      K   t         j                  |       } | D ]@  }t        |t              r|s|j	                  d      D ]  }|j
                  |k(  s|  B yw)am  Yields the `requires_grad` parameters of a given module or list of modules.

        Args:
            modules: A given module or an iterable of modules
            train_bn: Whether not to train the BatchNorm module
            requires_grad: Whether to create a generator for trainable or non-trainable parameters.
        Returns:
            Generator

        FrecurseN)r   r;   r8   r   
parametersrD   )r6   rC   rD   modparams        r   filter_paramszBaseFinetuning.filter_params   sY      !009C#z*86&&-7K 7	 s   AA	Ac                     t         j                  |       } | D ]7  }t        |t              rd|_        |j                  d      D ]	  }d|_         9 y)zUnfreezes the parameters of the provided modules.

        Args:
            modules: A given module or an iterable of modules

        TFrF   N)r   r;   r8   r   track_running_statsrH   rD   )r6   modulerJ   s      r   make_trainablezBaseFinetuning.make_trainable   sP     !009F&*--1***5*9&*# :	 r   rN   c                 n    t        | t              rd| _        | j                  d      D ]	  }d|_         y)zjFreezes the parameters of the provided module.

        Args:
            module: A given module

        FrF   N)r8   r   rM   rH   rD   )rN   rJ   s     r   freeze_modulezBaseFinetuning.freeze_module   s6     fj)).F&&&u&5E"'E 6r   c                     t         j                  |       } | D ]?  }t        |t              r|rt         j	                  |       +t         j                  |       A y)zFreezes the parameters of the provided modules.

        Args:
            modules: A given module or an iterable of modules
            train_bn: If True, leave the BatchNorm layers in training mode

        Returns:
            None

        N)r   r;   r8   r   rO   rQ   )r6   rC   rI   s      r   freezezBaseFinetuning.freeze   sF     !009C#z*x--c2,,S1	 r   r4   paramsc                     g }g }|D ]C  t        fd| j                  D              s|j                         3|j                         E |rt        dt	        |        d       |S )ac  This function is used to exclude any parameter which already exists in this optimizer.

        Args:
            optimizer: Optimizer used for parameter exclusion
            params: Iterable of parameters used to check against the provided optimizer

        Returns:
            List of parameters not contained in this optimizer param groups

        c              3   ^   K   | ]$  }|d    D ]  }t        j                  |        & yw)rT   N)torchequal).0groupprJ   s      r   	<genexpr>z5BaseFinetuning.filter_on_optimizer.<locals>.<genexpr>   s1     j:PZ_`hZiUVu{{1e,Zi,:Ps   *-zThe provided params to be frozen already exist within another group of this optimizer. Those parameters will be skipped.
HINT: Did you init your optimizer in `configure_optimizer` as such:
 z<(filter(lambda p: p.requires_grad, self.parameters()), ...) )anyr2   appendr   type)r4   rT   
out_paramsremoved_paramsrJ   s       @r   filter_on_optimizerz"BaseFinetuning.filter_on_optimizer   sx     
Ej):P:Pjj!!%(%%e,	   O$$`b r   lrinitial_denom_lrc                    t         j                  |        ||j                  d   d   n
t        |      }||nd}t         j	                  | |d      }t         j                  ||      }|r|j                  |||z  d       yy)a  Unfreezes a module and adds its parameters to an optimizer.

        Args:
            modules: A module or iterable of modules to unfreeze.
                Their parameters will be added to an optimizer as a new param group.
            optimizer: The provided optimizer will receive new parameters and will add them to
                `add_param_group`
            lr: Learning rate for the new param group.
            initial_denom_lr: If no lr is provided, the learning from the first param group will be used
                and divided by `initial_denom_lr`.
            train_bn: Whether to train the BatchNormalization layers.

        Nr   rc   g      ?T)rC   rD   )rT   rc   )r   rO   r2   floatrK   rb   add_param_group)r6   r4   rc   rd   rC   	params_lrdenom_lrrT   s           r   unfreeze_and_add_param_groupz+BaseFinetuning.unfreeze_and_add_param_group   s    * 	%%g.79zI**1-d3uRy	')z#s--gX\-]33IvF%%y8?S&TU r   stagec                 t    | j                  |       ddlm} t        |j                  |      rt        d      y )Nr   )DeepSpeedStrategyzThe Finetuning callback does not support running with the DeepSpeed strategy. Choose a different strategy or disable the callback.)freeze_before_trainingpytorch_lightning.strategiesrm   r8   strategyNotImplementedError)r!   r(   r*   rk   rm   s        r   setupzBaseFinetuning.setup  s<    ##I.Bg&&(9:%H  ;r   r2   mappingc                     g }| D ]U  }|j                         D ci c]  \  }}|dk7  s|| }}}|d   D cg c]  }||   	 c}|d<   |j                  |       W |S c c}}w c c}w )NrT   )itemsr^   )r2   rs   outputgkvgroup_stater[   s           r   r1   z-BaseFinetuning._apply_mapping_to_param_groups'  sv    A,-GGIGIDAqh1a4IKG9:8$EAWQZ$EK!MM+&	 
  H$Es   AAA%r3   num_param_groupscurrent_param_groupsc                 2   |j                         D ci c]  \  }}||
 }}}|| j                  vr | j                  ||      | j                  |<   y |t        |      k7  r2| j                  |   j	                  | j                  ||d  |             y y c c}}w N)r.   r   r1   lenr:   )r!   r*   r3   r{   r|   nr[   rs   s           r   _storezBaseFinetuning._store1  s     %.$>$>$@A$@DAq1a4$@A$;;;9=9\9\$g:D--g6 %9!::--g6==334HIYIZ4[]de ; Bs   Bc                     t        |j                        D ]W  \  }}t        |j                        }| j	                  ||j
                  |       |j                  }| j                  ||||       Y y)Called when the epoch begins.N)r/   r0   r   r2   finetune_functioncurrent_epochr   )r!   r(   r*   r3   r4   r{   r|   s          r   on_train_epoch_startz#BaseFinetuning.on_train_epoch_startC  sh     #,G,>,>"?GY"9#9#9:""9g.C.CYO#,#9#9 KK	7,<>RS	 #@r   r   c                     t         )z$Override to add your unfreeze logic.rq   )r!   r*   r   r4   s       r   r   z BaseFinetuning.finetune_functionL      !!r   c                     t         )z"Override to add your freeze logic.r   r!   r*   s     r   rn   z%BaseFinetuning.freeze_before_trainingP  r   r   )r   Nr(   r)   r*   r+   r   N)TT)T)N      $@Tr*   r+   r   N)$__name__
__module____qualname____doc__r"   r   r-   strr   r%   r'   r5   staticmethodr   r	   r   r=   r;   boolr   rK   rO   rQ   rS   r   rb   r   rf   rj   rr   r1   intr   r   r   rn   r   r   r   r   r   0   sI   *X! 
DcN 
 

 ;$sCx. ;T ; ; % % Pvxfh>N8O/P'P!Q PVZ[aVb P P4 pt vxfh.>(?@@A MQ im 	   , +fhuVX=M7N.O&O P +UY + + (f ( ( ( 2fhuVX5E/F&GGH 2TX 2dh 2 2$ y ( t  8  #"&Vvxfh.>(?@@AVV UOV  	V
 V 
V V8 	\ 	6J 	SV 	[_ 	 	 T$sCx.5I TX ]abfgjlogobp]q  '  	
 #4S>2 
$ T T"+? " "Xa "fj ""r   r   c                        e Zd ZdZdedddddddf	d	ed
ededee   de	dede	de	deddf fdZ
edeeef   fd       Zedeeef   ddf fd       Zed fd       Zedd       Zedddededdfd       Z xZS ) r   a^  Finetune a backbone model based on a learning rate user-defined scheduling.

    When the backbone learning rate reaches the current model learning rate
    and ``should_align`` is set to True, it will align with it for the rest of the training.

    Args:
        unfreeze_backbone_at_epoch: Epoch at which the backbone will be unfreezed.
        lambda_func: Scheduling function for increasing backbone learning rate.
        backbone_initial_ratio_lr:
            Used to scale down the backbone learning rate compared to rest of model
        backbone_initial_lr: Optional, Initial learning rate for the backbone.
            By default, we will use ``current_learning /  backbone_initial_ratio_lr``
        should_align: Whether to align with current learning rate when backbone learning
            reaches it.
        initial_denom_lr: When unfreezing the backbone, the initial learning rate will
            ``current_learning_rate /  initial_denom_lr``.
        train_bn: Whether to make Batch Normalization trainable.
        verbose: Display current learning rate for model and backbone
        rounding: Precision for displaying learning rate

    Example::

        >>> import torch
        >>> import torch.nn as nn
        >>> from pytorch_lightning import LightningModule, Trainer
        >>> from pytorch_lightning.callbacks import BackboneFinetuning
        >>> import torchvision.models as models
        >>>
        >>> class TransferLearningModel(LightningModule):
        ...     def __init__(self, num_classes=10):
        ...         super().__init__()
        ...         # REQUIRED: Your model must have a 'backbone' attribute
        ...         self.backbone = models.resnet50(weights=None)
        ...         # Remove the final classification layer from backbone
        ...         self.backbone = nn.Sequential(*list(self.backbone.children())[:-1])
        ...
        ...         # Add your task-specific head
        ...         self.head = nn.Sequential(
        ...             nn.Flatten(),
        ...             nn.Linear(2048, 512),
        ...             nn.ReLU(),
        ...             nn.Linear(512, num_classes)
        ...         )
        ...
        ...     def forward(self, x):
        ...         # Extract features with backbone
        ...         features = self.backbone(x)
        ...         # Classify with head
        ...         return self.head(features)
        ...
        ...     def configure_optimizers(self):
        ...         # Initially only optimize the head - backbone will be added by callback
        ...         return torch.optim.Adam(self.head.parameters(), lr=1e-3)
        ...
        >>> # Setup the callback
        >>> multiplicative = lambda epoch: 1.5
        >>> backbone_finetuning = BackboneFinetuning(
        ...     unfreeze_backbone_at_epoch=10,  # Start unfreezing at epoch 10
        ...     lambda_func=multiplicative,     # Gradually increase backbone LR
        ...     backbone_initial_ratio_lr=0.1,  # Start backbone at 10% of head LR
        ... )
        >>> model = TransferLearningModel()
        >>> trainer = Trainer(callbacks=[backbone_finetuning])

    
   g?NTr   F   unfreeze_backbone_at_epochlambda_funcbackbone_initial_ratio_lrbackbone_initial_lrshould_alignrd   rC   verboseroundingr   c
                     t         
|           || _        || _        || _        || _        || _        || _        || _        || _	        |	| _
        d | _        y r~   )superr"   r   r   r   r   r   rd   rC   r   r   previous_backbone_lr)r!   r   r   r   r   r   rd   rC   r   r   	__class__s             r   r"   zBackboneFinetuning.__init__  s]     	/I'%00I&4G ".'7&$%59!r   c                 4    | j                   | j                  dS )N)r$   r   )r   r   r    s    r   r%   zBackboneFinetuning.state_dict  s      ,0+L+L$($=$=
 	
r   r%   c                 8    |d   | _         t        | 	  |       y )Nr   )r   r   r'   )r!   r%   r   s     r   r'   z"BackboneFinetuning.load_state_dict  s    $./E$F!
+r   r*   r+   c                     t        |d      r*t        |j                  t              rt        |   ||      S t        d      )z
        Raises:
            MisconfigurationException:
                If LightningModule has no nn.Module `backbone` attribute.
        backbonez@The LightningModule should have a nn.Module `backbone` attribute)hasattrr8   r   r	   r   r5   r   )r!   r(   r*   r   s      r   r5   zBackboneFinetuning.on_fit_start  s<     9j)j9K9KV.T7';;'(jkkr   c                 :    | j                  |j                         y r~   )rS   r   r   s     r   rn   z)BackboneFinetuning.freeze_before_training  s    I&&'r   r   r4   c           	      .   || j                   k(  r|j                  d   d   }| j                  | j                  n|| j                  z  }|| _        | j                  |j                  ||| j                  | j                         | j                  rDt        j                  dt        || j                         dt        || j                                yy|| j                   kD  r|j                  d   d   }| j                  |dz         | j                  z  }| j                  r||kD  r|n|}||j                  d   d<   || _        | j                  rDt        j                  dt        || j                         dt        || j                                yyy)	r   r   rc   N)rC   rd   zCurrent lr: z, Backbone lr:    )r   r2   r   r   r   rj   r   rC   rd   r   loginforoundr   r   r   )r!   r*   r   r4   
current_lrinitial_backbone_lrnext_current_backbone_lrs          r   r   z$BackboneFinetuning.finetune_function  s    D333"//248J ++7 (($"@"@@  
 )<D%--""#!%!6!6 .  ||"5T]]#C"D E$$)*=t}}$M#NP  T444"//248J'+'7'7	'BTE^E^'^$ %%*BZ*O - %
 0HI""2&t,(@D%||"5T]]#C"D E$$)*BDMM$R#SU  5r   r   r   )r   r   r   r   r   r   r   rf   r   r   r"   r   r-   r   r   r%   r'   r5   rn   r   r   __classcell__)r   s   @r   r   r   U  s<   @H +- .+0/3!"&:$': : $)	:
 &e_: :  : : : : 
:2 
DcN 
 
 ,$sCx. ,T , , l l ( ( %+? % %Xa %fj % %r   )'r   loggingcollections.abcr   r   typingr   r   r   r   rW   torch.nnr	   r
   torch.nn.modules.batchnormr   torch.optim.optimizerr   typing_extensionsr   pytorch_lightningpl$pytorch_lightning.callbacks.callbackr   &pytorch_lightning.utilities.exceptionsr   #pytorch_lightning.utilities.importsr   %pytorch_lightning.utilities.rank_zeror   __doctest_skip__	getLoggerr   r   r   rf   r   r   r   r   r   r   <module>r      s     / 1 1  ' 1 + &  9 L F @,- g!# % b"X b"J	] ]r   