class_docstr: This class extends optimizers with Stochastic Weight Averaging (SWA).
  
  The Stochastic Weight Averaging mechanism was proposed by Pavel Izmailov
  et. al in the paper [Averaging Weights Leads to Wider Optima and
  Better Generalization](https://arxiv.org/abs/1803.05407). The optimizer
  implements averaging of multiple points along the trajectory of SGD. The
  optimizer expects an inner optimizer which will be used to apply the
  gradients to the variables and itself computes a running average of the
  variables every `k` steps (which generally corresponds to the end
  of a cycle when a cyclic learning rate is employed).
  
  We also allow the specification of the number of steps averaging
  should first happen after. Let's say, we want averaging to happen every `k`
  steps after the first `m` steps. After step `m` we'd take a snapshot of the
  variables and then average the weights appropriately at step `m + k`,
  `m + 2k` and so on. The assign_average_vars function can be called at the
  end of training to obtain the averaged_weights from the optimizer.
  
  Note: If your model has batch-normalization layers you would need to run
  the final weights through the data to compute the running mean and
  variance corresponding to the activations for each layer of the network.
  From the paper: If the DNN uses batch normalization we run one
  additional pass over the data, to compute the running mean and standard
  deviation of the activations for each layer of the network with SWA
  weights after the training is finished, since these statistics are not
  collected during training. For most deep learning libraries, such as
  PyTorch or Tensorflow, one can typically collect these statistics by
  making a forward pass over the data in training mode
  ([Averaging Weights Leads to Wider Optima and Better
  Generalization](https://arxiv.org/abs/1803.05407))
  
  Example of usage:
  
  ```python
  opt = tf.keras.optimizers.SGD(learning_rate)
  opt = tfa.optimizers.SWA(opt, start_averaging=m, average_period=k)
  ```
docstr: Wrap optimizer with the Stochastic Weight Averaging mechanism.
  
  Args:
      optimizer: The original optimizer that will be used to compute and
          apply the gradients.
      start_averaging: An integer. Threshold to start averaging using
          SWA. Averaging only occurs at `start_averaging` iters, must
          be >= 0. If start_averaging = m, the first snapshot will be
          taken after the mth application of gradients (where the first
          iteration is iteration 0).
      average_period: An integer. The synchronization period of SWA. The
          averaging occurs every average_period steps. Averaging period
          needs to be >= 1.
      name: Optional name for the operations created when applying
          gradients. Defaults to 'SWA'.
      **kwargs: keyword arguments. Allowed to be {`clipnorm`,
          `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients by
          norm; `clipvalue` is clip gradients by value, `decay` is
          included for backward compatibility to allow time inverse
          decay of learning rate. `lr` is included for backward
          compatibility, recommended to use `learning_rate` instead.
code: def __init__(
      self,
      optimizer: types.Optimizer,
      start_averaging: int = 0,
      average_period: int = 10,
      name: str = "SWA",
      **kwargs,
  ):
      r"""Wrap optimizer with the Stochastic Weight Averaging mechanism.
  
      Args:
          optimizer: The original optimizer that will be used to compute and
              apply the gradients.
          start_averaging: An integer. Threshold to start averaging using
              SWA. Averaging only occurs at `start_averaging` iters, must
              be >= 0. If start_averaging = m, the first snapshot will be
              taken after the mth application of gradients (where the first
              iteration is iteration 0).
          average_period: An integer. The synchronization period of SWA. The
              averaging occurs every average_period steps. Averaging period
              needs to be >= 1.
          name: Optional name for the operations created when applying
              gradients. Defaults to 'SWA'.
          **kwargs: keyword arguments. Allowed to be {`clipnorm`,
              `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients by
              norm; `clipvalue` is clip gradients by value, `decay` is
              included for backward compatibility to allow time inverse
              decay of learning rate. `lr` is included for backward
              compatibility, recommended to use `learning_rate` instead.
      """
      super().__init__(optimizer, name, **kwargs)
  
      if average_period < 1:
          raise ValueError("average_period must be >= 1")
      if start_averaging < 0:
          raise ValueError("start_averaging must be >= 0")
  
      self._set_hyper("average_period", average_period)
      self._set_hyper("start_averaging", start_averaging)
  
file_path: addons/tensorflow_addons/optimizers/stochastic_weight_averaging.py
imports: import tensorflow
  from tensorflow_addons.optimizers.average_wrapper import AveragedOptimizerWrapper
  from tensorflow_addons.utils import types
  from typeguard import typechecked
label_desc: Use image node_id 1 to create a new SWA object from inherited base classes: AveragedOptimizerWrapper with example: obj = SWA(optimizer, start_averaging, average_period, name)