docstr: Implements Stochastic Weight Averaging (SWA).
  
  THIS WAS DIRECTLY COPIED OVER FROM pytorch/contrib
      https://github.com/pytorch/contrib
  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).
  
  SWA is implemented as a wrapper class taking optimizer instance as input
  and applying SWA on top of that optimizer.
  
  SWA can be used in two modes: automatic and manual. In the automatic
  mode SWA running averages are automatically updated every
  :attr:`swa_freq` steps after :attr:`swa_start` steps of optimization. If
  :attr:`swa_lr` is provided, the learning rate of the optimizer is reset
  to :attr:`swa_lr` at every step starting from :attr:`swa_start`. To use
  SWA in automatic mode provide values for both :attr:`swa_start` and
  :attr:`swa_freq` arguments.
  
  Alternatively, in the manual mode, use :meth:`update_swa` or
  :meth:`update_swa_group` methods to update the SWA running averages.
  
  In the end of training use `swap_swa_sgd` method to set the optimized
  variables to the computed averages.
  
  Args:
      optimizer (torch.optim.Optimizer): optimizer to use with SWA
      swa_start (int): number of steps before starting to apply SWA in
          automatic mode; if None, manual mode is selected (default: None)
      swa_freq (int): number of steps between subsequent updates of
          SWA running averages in automatic mode; if None, manual mode is
          selected (default: None)
      swa_lr (float): learning rate to use starting from step swa_start
          in automatic mode; if None, learning rate is not changed
          (default: None)
  
  .. _Averaging Weights Leads to Wider Optima and Better Generalization:
      https://arxiv.org/abs/1803.05407
  .. _Improving Consistency-Based Semi-Supervised Learning with Weight
      Averaging:
      https://arxiv.org/abs/1806.05594
docstr_tok: ["Implements","Stochastic","Weight","Averaging","(","SWA",")",".","THIS","WAS","DIRECTLY","COPIED","OVER","FROM","pytorch\/contrib","https",":","\/\/github.com\/pytorch\/contrib","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",")",".","SWA","is","implemented","as","a","wrapper","class","taking","optimizer","instance","as","input","and","applying","SWA","on","top","of","that","optimizer",".","SWA","can","be","used","in","two","modes",":","automatic","and","manual",".","In","the","automatic","mode","SWA","running","averages","are","automatically","updated","every",":","attr",":","`","swa_freq","`","steps","after",":","attr",":","`","swa_start","`","steps","of","optimization",".","If",":","attr",":","`","swa_lr","`","is","provided",",","the","learning","rate","of","the","optimizer","is","reset","to",":","attr",":","`","swa_lr","`","at","every","step","starting","from",":","attr",":","`","swa_start","`",".","To","use","SWA","in","automatic","mode","provide","values","for","both",":","attr",":","`","swa_start","`","and",":","attr",":","`","swa_freq","`","arguments",".","Alternatively",",","in","the","manual","mode",",","use",":","meth",":","`","update_swa","`","or",":","meth",":","`","update_swa_group","`","methods","to","update","the","SWA","running","averages",".","In","the","end","of","training","use","`","swap_swa_sgd","`","method","to","set","the","optimized","variables","to","the","computed","averages",".","Args",":","optimizer","(","torch.optim.Optimizer",")",":","optimizer","to","use","with","SWA","swa_start","(","int",")",":","number","of","steps","before","starting","to","apply","SWA","in","automatic","mode",";","if","None",",","manual","mode","is","selected","(","default",":","None",")","swa_freq","(","int",")",":","number","of","steps","between","subsequent","updates","of","SWA","running","averages","in","automatic","mode",";","if","None",",","manual","mode","is","selected","(","default",":","None",")","swa_lr","(","float",")",":","learning","rate","to","use","starting","from","step","swa_start","in","automatic","mode",";","if","None",",","learning","rate","is","not","changed","(","default",":","None",")","..","_Averaging","Weights","Leads","to","Wider","Optima","and","Better","Generalization",":","https",":","\/\/arxiv.org\/abs\/1803.05407","..","_Improving","Consistency-Based","Semi-Supervised","Learning","with","Weight","Averaging",":","https",":","\/\/arxiv.org\/abs\/1806.05594"]
code: def __init__(
      self, optimizer, swa_start=None, swa_freq=None, swa_lr=None
  ):
      r"""Implements Stochastic Weight Averaging (SWA).
  
      THIS WAS DIRECTLY COPIED OVER FROM pytorch/contrib
          https://github.com/pytorch/contrib
      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).
  
      SWA is implemented as a wrapper class taking optimizer instance as input
      and applying SWA on top of that optimizer.
  
      SWA can be used in two modes: automatic and manual. In the automatic
      mode SWA running averages are automatically updated every
      :attr:`swa_freq` steps after :attr:`swa_start` steps of optimization. If
      :attr:`swa_lr` is provided, the learning rate of the optimizer is reset
      to :attr:`swa_lr` at every step starting from :attr:`swa_start`. To use
      SWA in automatic mode provide values for both :attr:`swa_start` and
      :attr:`swa_freq` arguments.
  
      Alternatively, in the manual mode, use :meth:`update_swa` or
      :meth:`update_swa_group` methods to update the SWA running averages.
  
      In the end of training use `swap_swa_sgd` method to set the optimized
      variables to the computed averages.
  
      Args:
          optimizer (torch.optim.Optimizer): optimizer to use with SWA
          swa_start (int): number of steps before starting to apply SWA in
              automatic mode; if None, manual mode is selected (default: None)
          swa_freq (int): number of steps between subsequent updates of
              SWA running averages in automatic mode; if None, manual mode is
              selected (default: None)
          swa_lr (float): learning rate to use starting from step swa_start
              in automatic mode; if None, learning rate is not changed
              (default: None)
  
      .. _Averaging Weights Leads to Wider Optima and Better Generalization:
          https://arxiv.org/abs/1803.05407
      .. _Improving Consistency-Based Semi-Supervised Learning with Weight
          Averaging:
          https://arxiv.org/abs/1806.05594
      """
      self._auto_mode, (
          self.swa_start,
          self.swa_freq,
      ) = self._check_params(self, swa_start, swa_freq)
      self.swa_lr = swa_lr
  
      if self._auto_mode:
          if swa_start < 0:
              raise ValueError(
                  "Invalid swa_start: {}".format(swa_start)
              )
          if swa_freq < 1:
              raise ValueError("Invalid swa_freq: {}".format(swa_freq))
      else:
          if self.swa_lr is not None:
              warnings.warn(
                  "Some of swa_start, swa_freq is None, ignoring swa_lr"
              )
          # If not in auto mode make all swa parameters None
          self.swa_lr = None
          self.swa_start = None
          self.swa_freq = None
  
      if self.swa_lr is not None and self.swa_lr < 0:
          raise ValueError(
              "Invalid SWA learning rate: {}".format(swa_lr)
          )
  
      self.optimizer = optimizer
  
      self.defaults = self.optimizer.defaults
      self.param_groups = self.optimizer.param_groups
      self.state = defaultdict(dict)
      self.opt_state = self.optimizer.state
      for group in self.param_groups:
          group["n_avg"] = 0
          group["step_counter"] = 0
  
imports: import warnings
  from collections import defaultdict
  from typing import Optional, Union
  import torch
  from pytext.config.component import create_optimizer
  from pytext.optimizer.lamb import Lamb
  from pytext.optimizer.madgrad import MADGRAD
  from pytext.optimizer.optimizers import Adagrad, Adam, AdamW, Optimizer, SGD
  from pytext.optimizer.radam import RAdam
  from torch.optim import Optimizer
label_desc: Use image node_id 1 to create a new StochasticWeightAveraging object from inherited base classes: Optimizer, PT_Optimizer with example: obj = StochasticWeightAveraging(optimizer, swa_start, swa_freq, swa_lr)