docstr: Implements the Stochastic Weight Averaging (SWA) Callback to average a model.
  
  Stochastic Weight Averaging was proposed in ``Averaging Weights Leads to
  Wider Optima and Better Generalization`` by Pavel Izmailov, Dmitrii
  Podoprikhin, Timur Garipov, Dmitry Vetrov and Andrew Gordon Wilson
  (UAI 2018).
  
  This documentation is highly inspired by PyTorch's work on SWA.
  The callback arguments follow the scheme defined in PyTorch's ``swa_utils`` package.
  
  For a SWA explanation, please take a look
  `here <https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging>`_.
  
  .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
  
  .. warning:: ``StochasticWeightAveraging`` is currently not supported for multiple optimizers/schedulers.
  
  .. warning:: ``StochasticWeightAveraging`` is currently only supported on every epoch.
  
  See also how to :ref:`enable it directly on the Trainer <advanced/training_tricks:Stochastic Weight Averaging>`
  
  Arguments:
  
      swa_lrs: The SWA learning rate to use:
  
          - ``float``. Use this value for all parameter groups of the optimizer.
          - ``List[float]``. A list values for each parameter group of the optimizer.
  
      swa_epoch_start: If provided as int, the procedure will start from
          the ``swa_epoch_start``-th epoch. If provided as float between 0 and 1,
          the procedure will start from ``int(swa_epoch_start * max_epochs)`` epoch
  
      annealing_epochs: number of epochs in the annealing phase (default: 10)
  
      annealing_strategy: Specifies the annealing strategy (default: "cos"):
  
          - ``"cos"``. For cosine annealing.
          - ``"linear"`` For linear annealing
  
      avg_fn: the averaging function used to update the parameters;
          the function must take in the current value of the
          :class:`AveragedModel` parameter, the current value of :attr:`model`
          parameter and the number of models already averaged; if None,
          equally weighted average is used (default: ``None``)
  
      device: if provided, the averaged model will be stored on the ``device``.
          When None is provided, it will infer the `device` from ``pl_module``.
          (default: ``"cpu"``)
docstr_tok: ["Implements","the","Stochastic","Weight","Averaging","(","SWA",")","Callback","to","average","a","model",".","Stochastic","Weight","Averaging","was","proposed","in","``","Averaging","Weights","Leads","to","Wider","Optima","and","Better","Generalization","``","by","Pavel","Izmailov",",","Dmitrii","Podoprikhin",",","Timur","Garipov",",","Dmitry","Vetrov","and","Andrew","Gordon","Wilson","(","UAI","2018",")",".","This","documentation","is","highly","inspired","by","PyTorch","'s","work","on","SWA",".","The","callback","arguments","follow","the","scheme","defined","in","PyTorch","'s","``","swa_utils","``","package",".","For","a","SWA","explanation",",","please","take","a","look","`","here","<","https",":","\/\/pytorch.org\/blog\/pytorch-1.6-now-includes-stochastic-weight-averaging",">","`","_",".","..","warning",":",":","This","is","an",":","ref",":","`","experimental","<","versioning",":","Experimental","API",">","`","feature",".","..","warning",":",":","``","StochasticWeightAveraging","``","is","currently","not","supported","for","multiple","optimizers\/schedulers",".","..","warning",":",":","``","StochasticWeightAveraging","``","is","currently","only","supported","on","every","epoch",".","See","also","how","to",":","ref",":","`","enable","it","directly","on","the","Trainer","<","advanced\/training_tricks",":","Stochastic","Weight","Averaging",">","`","Arguments",":","swa_lrs",":","The","SWA","learning","rate","to","use",":","-","``","float","``",".","Use","this","value","for","all","parameter","groups","of","the","optimizer",".","-","``","List","[","float","]","``",".","A","list","values","for","each","parameter","group","of","the","optimizer",".","swa_epoch_start",":","If","provided","as","int",",","the","procedure","will","start","from","the","``","swa_epoch_start","``","-th","epoch",".","If","provided","as","float","between","0","and","1",",","the","procedure","will","start","from","``","int","(","swa_epoch_start","*","max_epochs",")","``","epoch","annealing_epochs",":","number","of","epochs","in","the","annealing","phase","(","default",":","10",")","annealing_strategy",":","Specifies","the","annealing","strategy","(","default",":","``","cos","''",")",":","-","``","``","cos","''","``",".","For","cosine","annealing",".","-","``","``","linear","''","``","For","linear","annealing","avg_fn",":","the","averaging","function","used","to","update","the","parameters",";","the","function","must","take","in","the","current","value","of","the",":","class",":","`","AveragedModel","`","parameter",",","the","current","value","of",":","attr",":","`","model","`","parameter","and","the","number","of","models","already","averaged",";","if","None",",","equally","weighted","average","is","used","(","default",":","``","None","``",")","device",":","if","provided",",","the","averaged","model","will","be","stored","on","the","``","device","``",".","When","None","is","provided",",","it","will","infer","the","`","device","`","from","``","pl_module","``",".","(","default",":","``","``","cpu","''","``",")"]
code: def __init__(
      self,
      swa_lrs: Union[float, List[float]],
      swa_epoch_start: Union[int, float] = 0.8,
      annealing_epochs: int = 10,
      annealing_strategy: str = "cos",
      avg_fn: Optional[_AVG_FN] = None,
      device: Optional[Union[torch.device, str]] = torch.device("cpu"),
  ):
      r"""Implements the Stochastic Weight Averaging (SWA) Callback to average a model.
  
      Stochastic Weight Averaging was proposed in ``Averaging Weights Leads to
      Wider Optima and Better Generalization`` by Pavel Izmailov, Dmitrii
      Podoprikhin, Timur Garipov, Dmitry Vetrov and Andrew Gordon Wilson
      (UAI 2018).
  
      This documentation is highly inspired by PyTorch's work on SWA.
      The callback arguments follow the scheme defined in PyTorch's ``swa_utils`` package.
  
      For a SWA explanation, please take a look
      `here <https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging>`_.
  
      .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
  
      .. warning:: ``StochasticWeightAveraging`` is currently not supported for multiple optimizers/schedulers.
  
      .. warning:: ``StochasticWeightAveraging`` is currently only supported on every epoch.
  
      See also how to :ref:`enable it directly on the Trainer <advanced/training_tricks:Stochastic Weight Averaging>`
  
      Arguments:
  
          swa_lrs: The SWA learning rate to use:
  
              - ``float``. Use this value for all parameter groups of the optimizer.
              - ``List[float]``. A list values for each parameter group of the optimizer.
  
          swa_epoch_start: If provided as int, the procedure will start from
              the ``swa_epoch_start``-th epoch. If provided as float between 0 and 1,
              the procedure will start from ``int(swa_epoch_start * max_epochs)`` epoch
  
          annealing_epochs: number of epochs in the annealing phase (default: 10)
  
          annealing_strategy: Specifies the annealing strategy (default: "cos"):
  
              - ``"cos"``. For cosine annealing.
              - ``"linear"`` For linear annealing
  
          avg_fn: the averaging function used to update the parameters;
              the function must take in the current value of the
              :class:`AveragedModel` parameter, the current value of :attr:`model`
              parameter and the number of models already averaged; if None,
              equally weighted average is used (default: ``None``)
  
          device: if provided, the averaged model will be stored on the ``device``.
              When None is provided, it will infer the `device` from ``pl_module``.
              (default: ``"cpu"``)
  
      """
  
      err_msg = "swa_epoch_start should be a >0 integer or a float between 0 and 1."
      if isinstance(swa_epoch_start, int) and swa_epoch_start < 1:
          raise MisconfigurationException(err_msg)
      if isinstance(swa_epoch_start, float) and not (
          0 <= swa_epoch_start <= 1
      ):
          raise MisconfigurationException(err_msg)
  
      wrong_type = not isinstance(swa_lrs, (float, list))
      wrong_float = isinstance(swa_lrs, float) and swa_lrs <= 0
      wrong_list = isinstance(swa_lrs, list) and not all(
          lr > 0 and isinstance(lr, float) for lr in swa_lrs
      )
      if wrong_type or wrong_float or wrong_list:
          raise MisconfigurationException(
              "The `swa_lrs` should a positive float, or a list of positive floats"
          )
  
      if avg_fn is not None and not callable(avg_fn):
          raise MisconfigurationException(
              "The `avg_fn` should be callable."
          )
  
      if device is not None and not isinstance(
          device, (torch.device, str)
      ):
          raise MisconfigurationException(
              f"device is expected to be a torch.device or a str. Found {device}"
          )
  
      self.n_averaged: Optional[Tensor] = None
      self._swa_epoch_start = swa_epoch_start
      self._swa_lrs = swa_lrs
      self._annealing_epochs = annealing_epochs
      self._annealing_strategy = annealing_strategy
      self._avg_fn = avg_fn or self.avg_fn
      self._device = device
      self._model_contains_batch_norm: Optional[bool] = None
      self._average_model: Optional["pl.LightningModule"] = None
      self._initialized = False
      self._swa_scheduler: Optional[LRScheduler] = None
      self._scheduler_state: Optional[Dict] = None
      self._init_n_averaged = 0
      self._latest_update_epoch = -1
      self.momenta: Dict[
          nn.modules.batchnorm._BatchNorm, Optional[float]
      ] = {}
      self._max_epochs: int
  
imports: from copy import deepcopy
  from typing import Any, Callable, Dict, List, Optional, Union, cast
  import torch
  from torch import Tensor, nn
  from torch.optim.swa_utils import SWALR
  from typing_extensions import override
  import lightning.pytorch
  from lightning.fabric.utilities.types import LRScheduler
  from lightning.pytorch.callbacks.callback import Callback
  from lightning.pytorch.strategies import DeepSpeedStrategy
  from lightning.pytorch.strategies.fsdp import FSDPStrategy
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from lightning.pytorch.utilities.rank_zero import rank_zero_info, rank_zero_warn
  from lightning.pytorch.utilities.types import LRSchedulerConfig
label_desc: Use image node_id 1 to create a new StochasticWeightAveraging object from inherited base classes: Callback with example: obj = StochasticWeightAveraging(swa_lrs, swa_epoch_start, annealing_epochs, annealing_strategy, avg_fn, device)