# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/optimizer.finetune.ipynb.

# %% ../../nbs/optimizer.finetune.ipynb 1
# Contains code from:
# fastai - Apache License 2.0 - Copyright (c) 2023 fast.ai

# %% ../../nbs/optimizer.finetune.ipynb 4
from __future__ import annotations

from fastcore.basics import GetAttr

from fastai.optimizer import (Optimizer, _update, weight_decay, l2_reg, average_grad, sgd_step,
                              momentum_step, average_sqr_grad, rms_prop_step, step_stat, adam_step,
                              radam_step, qhadam_step, larc_layer_lr, larc_step, lamb_step)
from ..imports import *

# %% auto 0
__all__ = ['FineTuneOpt', 'fine_tune_wd', 'SGDFT', 'RMSPropFT', 'AdamFT', 'RAdamFT', 'QHAdamFT', 'LarcFT', 'LambFT',
           'LookaheadFT', 'rangerFT']

# %% ../../nbs/optimizer.finetune.ipynb 8
class FineTuneOpt(Optimizer):
    """
    Modification of the base optimizer class for the fastai library, updating `params` with `cbs`

    In combination with the `fine_tune_wd` callback, adds optional weight decay `ft_wd` towards the starting value, 
    to prevent overfitting to the new dataset during fine-tuning.

    By default, will not apply to the fine-tuning head, just the pretrained body.

    From: https://gist.github.com/crowsonkb/f646976de8033b371ea17cb9b1c1561f
    """
    _keep_on_clear = ['force_train', 'do_wd']
    def __init__(self,
        params:Tensor, # Parameters and hyper parameters
        cbs:list, # `Optimizer` callbacks
        train_bn:bool=True, # Batch normalization is always trained
        wd_ft_head:bool=False, # Apply fine tuning weight decay to model head
        **defaults # Default values to set on hyper parameters
    ):
        super().__init__(params, cbs, train_bn, **defaults)
        self.wd_ft_head, self.set_orig_p = wd_ft_head, True

    @torch.no_grad()
    def step(self, closure=None):
        if self.set_orig_p:
            self.set_orig_p = False
            n = slice(None) if self.wd_ft_head or len(self.param_lists)<=1 else slice(None, -1)
            for p,pg,state,hyper in self.all_params(n):
                state['orig_p'] = p.detach().clone()
                self.state[p] = state
        super().step(closure)

    def clear_state(self):
        self.set_orig_p = True
        super().clear_state()

    def state_dict(self):
        state = super().state_dict()
        state.update({'set_orig_p': self.set_orig_p})
        return state

    def load_state_dict(self, sd):
        self.set_orig_p = sd.pop('set_orig_p')
        super().load_state_dict(sd)

# %% ../../nbs/optimizer.finetune.ipynb 12
def fine_tune_wd(p, lr, ft_wd, orig_p=None, do_wd=True, **kwargs):
    "Weight decay `p` towards the starting value `orig_p`"
    if do_wd and ft_wd !=0 and orig_p is not None:
        p.lerp_(orig_p, lr*ft_wd)

# %% ../../nbs/optimizer.finetune.ipynb 15
def SGDFT(params, lr, mom=0., wd=0., ft_wd=0., decouple_wd=True, wd_ft_head=False):
    "A `Optimizer` for SGD with `lr` and `mom` and `params`"
    cbs = [weight_decay] if decouple_wd else [l2_reg]
    cbs += [fine_tune_wd]
    if mom != 0: cbs.append(average_grad)
    cbs.append(sgd_step if mom==0 else momentum_step)
    return FineTuneOpt(params, cbs, lr=lr, mom=mom, wd=wd, ft_wd=ft_wd, wd_ft_head=wd_ft_head)

# %% ../../nbs/optimizer.finetune.ipynb 19
def RMSPropFT(params, lr, sqr_mom=0.99, mom=0., wd=0., ft_wd=0., decouple_wd=True, wd_ft_head=False):
    "A `FineTuneOpt` for RMSProp with `lr`, `sqr_mom`, `mom` and `params`"
    cbs = [weight_decay] if decouple_wd else [l2_reg]
    cbs += [fine_tune_wd] + [average_sqr_grad] if mom==0. else [average_grad, average_sqr_grad]
    cbs.append(rms_prop_step)
    return FineTuneOpt(params, cbs, lr=lr, mom=mom, sqr_mom=sqr_mom, wd=wd, ft_wd=ft_wd, wd_ft_head=wd_ft_head)

# %% ../../nbs/optimizer.finetune.ipynb 21
def AdamFT(params, lr, mom=0.9, sqr_mom=0.99, eps=1e-5, wd=0.01, ft_wd=0., decouple_wd=True, wd_ft_head=False):
    "A `FineTuneOpt` for Adam with `lr`, `mom`, `sqr_mom`, `eps` and `params`"
    cbs = [weight_decay] if decouple_wd else [l2_reg]
    cbs += [fine_tune_wd, partial(average_grad, dampening=True), average_sqr_grad, step_stat, adam_step]
    return FineTuneOpt(params, cbs, lr=lr, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd, ft_wd=ft_wd, wd_ft_head=wd_ft_head)

# %% ../../nbs/optimizer.finetune.ipynb 23
def RAdamFT(params, lr, mom=0.9, sqr_mom=0.99, eps=1e-5, wd=0., ft_wd=0., beta=0., decouple_wd=True, wd_ft_head=False):
    "A `FineTuneOpt` for Adam with `lr`, `mom`, `sqr_mom`, `eps` and `params`"
    cbs = [weight_decay] if decouple_wd else [l2_reg]
    cbs += [fine_tune_wd, partial(average_grad, dampening=True), average_sqr_grad, step_stat, radam_step]
    return FineTuneOpt(params, cbs, lr=lr, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd, ft_wd=ft_wd, beta=beta, wd_ft_head=wd_ft_head)

# %% ../../nbs/optimizer.finetune.ipynb 25
def QHAdamFT(params, lr, mom=0.999, sqr_mom=0.999, nu_1=0.7, nu_2 = 1.0, eps=1e-8, wd=0., ft_wd=0., decouple_wd=True, wd_ft_head=False):
    "An `FineTuneOpt` for Adam with `lr`, `mom`, `sqr_mom`, `nus`, eps` and `params`"
    cbs = [weight_decay] if decouple_wd else [l2_reg]
    cbs += [fine_tune_wd, partial(average_grad, dampening=True), partial(average_sqr_grad, dampening=True), step_stat, qhadam_step]
    return FineTuneOpt(params, cbs, lr=lr, nu_1=nu_1, nu_2=nu_2 ,
                       mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd,
                       ft_wd=ft_wd, wd_ft_head=wd_ft_head)

# %% ../../nbs/optimizer.finetune.ipynb 27
def LarcFT(params, lr, mom=0.9, clip=True, trust_coeff=0.02, eps=1e-8, wd=0., ft_wd=0., decouple_wd=True, wd_ft_head=False):
    "A `FineTuneOpt` for Adam with `lr`, `mom`, `sqr_mom`, `eps` and `params`"
    cbs = [weight_decay] if decouple_wd else [l2_reg]
    cbs += [fine_tune_wd]
    if mom!=0.: cbs.append(average_grad)
    cbs += [partial(larc_layer_lr, clip=clip), larc_step]
    return FineTuneOpt(params, cbs, lr=lr, mom=mom, trust_coeff=trust_coeff, eps=eps, wd=wd, ft_wd=ft_wd, wd_ft_head=wd_ft_head)

# %% ../../nbs/optimizer.finetune.ipynb 29
def LambFT(params, lr, mom=0.9, sqr_mom=0.99, eps=1e-5, wd=0., ft_wd=0., decouple_wd=True, wd_ft_head=False):
    "A `FineTuneOpt` for Adam with `lr`, `mom`, `sqr_mom`, `eps` and `params`"
    cbs = [weight_decay] if decouple_wd else [l2_reg]
    cbs += [fine_tune_wd, partial(average_grad, dampening=True), average_sqr_grad, step_stat, lamb_step]
    return Optimizer(params, cbs, lr=lr, mom=mom, sqr_mom=sqr_mom, eps=eps, wd=wd, ft_wd=ft_wd, wd_ft_head=wd_ft_head)

# %% ../../nbs/optimizer.finetune.ipynb 31
class LookaheadFT(FineTuneOpt, GetAttr):
    "Wrap a `FineTuneOpt` `opt` in a Lookahead optimizer"
    _default='opt'
    def __init__(self, opt, k=6, alpha=0.5):
        store_attr('opt,k,alpha')
        self._init_state()

    def step(self, closure=None):
        if closure is not None: raise NotImplementedError("fastai optimizers currently do not support closure")
        if self.slow_weights is None: self._copy_weights()
        self.opt.step()
        self.count += 1
        if self.count%self.k != 0: return
        for slow_pg,fast_pg in zip(self.slow_weights,self.param_lists):
            for slow_p,fast_p in zip(slow_pg,fast_pg):
                slow_p.data.add_(fast_p.data-slow_p.data, alpha=self.alpha)
                fast_p.data.copy_(slow_p.data)

    def clear_state(self):
        self.opt.clear_state()
        self._init_state()

    def state_dict(self):
        state = self.opt.state_dict()
        state.update({'count': self.count, 'slow_weights': self.slow_weights})
        return state

    def load_state_dict(self, sd):
        self.count = sd.pop('count')
        self.slow_weights = sd.pop('slow_weights')
        self.opt.load_state_dict(sd)

    def _init_state(self): self.count,self.slow_weights = 0,None
    def _copy_weights(self): self.slow_weights = L(L(p.clone().detach() for p in pg) for pg in self.param_lists)

    @property
    def param_lists(self): return self.opt.param_lists
    @param_lists.setter
    def param_lists(self, v): self.opt.param_lists = v

# %% ../../nbs/optimizer.finetune.ipynb 33
@delegates(RAdamFT)
def rangerFT(p, lr, mom=0.95, wd=0.01, ft_wd=0., eps=1e-6, **kwargs):
    "Convenience method for `LookaheadFT` with `RAdamFT`"
    return LookaheadFT(RAdamFT(p, lr=lr, mom=mom, wd=wd, ft_wd=ft_wd, eps=eps, **kwargs))
