# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/callback.ema.ipynb.

# %% ../../nbs/callback.ema.ipynb 1
# EMA Callbacks are inspired by timm's ModelEmaV2: https://github.com/rwightman/pytorch-image-models/blob/main/timm/utils/model_ema.py
# PyTorch Image Models - Apache License 2.0 - Copyright (c) 2020 Ross Wightman

from __future__ import annotations

from copy import deepcopy

from fastai.callback.core import Callback
from fastai.callback.fp16 import MixedPrecision
from fastai.callback.schedule import SchedCos, _Annealer

from ..imports import *

# %% auto 0
__all__ = ['EMACallback', 'EMAWarmupCallback']

# %% ../../nbs/callback.ema.ipynb 5
class EMACallback(Callback):
    "Exponential Moving Average (EMA) of model weights with a fused update step"
    order,run_valid = MixedPrecision.order+1,False
    def __init__(self, 
        decay:float=0.9998, # EMA decay value
        start_epoch:Number=0, # Epoch to start EMA in percent of training steps (float) or epochs (int, index 0)
        ema_device:torch.device|str|None=None, # Device to store EMA weights. Defaults to model device
        validate_ema:bool=True, # Run validation metrics using EMA weights instead of model weights. If true, `ema_device` must match model device
        replace_weights:bool=False, # Replace model weights with EMA weights when finished training. If false, set `Learner.model_ema` to EMA weights
        foreach:bool|None=None, # Fuse EMA update step with PyTorch ForEach methods or use a standard for loop. Defaults to true if PyTorch 1.12+ and Cuda device detected
        resume:bool=False, # Resume from EMA weights from previous training saved to `Learner.model_ema`
    ):
        store_attr()
        self._do_ema = False
        if self.foreach is None and ema_device is None:
            self.foreach = ismin_torch('1.12') and torch.cuda.is_available()

        if self.foreach:
            self.foreach_decay = 1-self.decay
            if notmax_torch('1.12'):
                warn('EMACallback with foreach=True is untested on PyTorch {torch.__verson__}, recommended to use 1.12 or newer')

        if resume and self.start_epoch > 0:
            warn(f'Resuming from prior EMA weights but delaying EMA until start_epoch={start_epoch}')

    @torch.no_grad()
    def before_fit(self):
        if hasattr(self.learn, 'lr_finder') or hasattr(self.learn, "gather_preds"):
            self.run = False
            return

        if self.start_epoch < 1:
            self.start_epoch = int(self.start_epoch*self.n_epoch)

        if self.resume:
            self.ema_model = self.learn.model_ema.eval()
        else:
            self.ema_model = deepcopy(self.learn.model).eval()

        model_device = next(self.learn.model.parameters()).device
        self.ema_model.to(self.ema_device if self.ema_device is not None else model_device)
        ema_device = next(self.ema_model.parameters()).device

        self.model_tensors, self.ema_tensors = [], []
        for mt, et in zip(self.learn.model.parameters(), self.ema_model.parameters()):
            self.model_tensors.append(mt)
            self.ema_tensors.append(et)

        self._validate_ema = model_device == ema_device if self.validate_ema else False
        if self.foreach:
            assert model_device == ema_device, "EMA device must equal model device if using foreach"

    @torch.no_grad()
    def before_train(self):
        if self.epoch == self.start_epoch:
            self._do_ema = True
            if self.start_epoch > 0 and not self.resume:
                self.ema_model.load_state_dict(self.learn.model.state_dict())

    @torch.no_grad()
    def after_batch(self):
        if self._do_ema:
            if self.foreach:
                torch._foreach_sub_(self.ema_tensors, self.ema_tensors, alpha=self.foreach_decay)
                torch._foreach_add_(self.ema_tensors, self.model_tensors, alpha=self.foreach_decay)
            else:
                for mt, et in zip(self.model_tensors, self.ema_tensors):
                    et.copy_(self.decay * et + (1. - self.decay) * mt)

    @torch.no_grad()
    def before_validate(self):
        if self._do_ema and self._validate_ema:
            self.temp_model = self.learn.model
            self.learn.model = self.ema_model

    @torch.no_grad()
    def after_validate(self):
        if self._do_ema and self._validate_ema:
            self.learn.model = self.temp_model

    def after_fit(self):
        if self.replace_weights:
            self.learn.model = self.ema_model
            self.ema_model = None
        else:
            self.learn.model_ema = self.ema_model

# %% ../../nbs/callback.ema.ipynb 8
class EMAWarmupCallback(EMACallback):
    "Exponential Moving Average (EMA) of model weights with a warmup schedule and fused update step"
    order,run_valid = MixedPrecision.order+1,False
    def __init__(self,
        start_decay:float=0.9, # Initial EMA decay value
        final_decay:float=0.9998, # Final EMA decay value
        start_epoch:Number=0, # Epoch to start EMA warmup in percent of training steps (float) or epochs (int, index 0)
        final_epoch:Number=0.3, # Epoch to finish EMA warmup in percent of training steps (float) or epochs (int, index 0)
        schedule:Callable[..., _Annealer]=SchedCos, # EMA decay warmup schedule
        ema_device:torch.device|str|None=None, # Device to store EMA weights. Defaults to model device
        validate_ema:bool=True, # Run validation metrics using EMA weights instead of model weights. If true, `ema_device` must match model device
        replace_weights:bool=False, # Replace model weights with EMA weights when finished training. If false, set `Learner.model_ema` to EMA weights
        foreach:bool|None=None, # Fuse EMA update step with PyTorch ForEach methods or use a standard for loop. Defaults to true if PyTorch 1.12+ and Cuda device detected
        resume:bool=False, # Resume from EMA weights from previous training saved to `Learner.model_ema`
        logger_callback:str='wandb', # Log EMA decay to `logger_callback` using `Callback.name` if available
    ):
        super().__init__(decay=final_decay, start_epoch=start_epoch, ema_device=ema_device, 
                         validate_ema=validate_ema, replace_weights=replace_weights, 
                         foreach=foreach, resume=resume)
        store_attr()
        self.schedule = schedule(start_decay, final_decay)

    def before_fit(self):
        super().before_fit()

        if self.final_epoch < 1:
            self.final_epoch = int(self.final_epoch*self.n_epoch)

        self.warmup_pct = 0.
        self.warmup_epochs = self.final_epoch - self.start_epoch

        if self.warmup_epochs == 0:
            warn(f'start_epoch {self.start_epoch} is equal to final_epoch {self.final_epoch} which negates warmup')

        self._log_ema_decay = getattr(self, f'_{self.logger_callback}_log_ema_decay', noop)
        self.has_logger = hasattr(self.learn, self.logger_callback) and self._log_ema_decay != noop

    def after_batch(self):
        if self._do_ema:
            if self.epoch <= self.final_epoch:
                if self.epoch == self.final_epoch:
                    self.decay = self.final_decay
                else:
                    self.decay = self.schedule(self.warmup_pct)
                    self.warmup_pct += 1./(self.learn.n_iter*self.warmup_epochs)
                if self.foreach: 
                    self.foreach_decay = 1-self.decay

            super().after_batch()

        if self.has_logger:
            if self._do_ema:
                self._log_ema_decay(self.decay)
            else:
                self._log_ema_decay(0.)

# %% ../../nbs/callback.ema.ipynb 11
try:
    import wandb

    @patch
    def _wandb_log_ema_decay(self:EMAWarmupCallback, decay:float):
        wandb.log({'ema_decay': decay}, self.learn.wandb._wandb_step+1)
except:
    pass
