class_docstr: Hook which keeps a track of the exponential moving average (EMA) of the model's
  parameters and applies the EMA params to the model during the test phases.
  
  Saving the state in cpu will save gpu memory, but will make training slower since
  the model parameters will need to be moved to cpu before the averaging.
  
  Note:
      This hooks stores two additional copies of the model's parameters, which will
      increase memory usage significantly.
class_docstr_tok: ["Hook","which","keeps","a","track","of","the","exponential","moving","average","(","EMA",")","of","the","model's","parameters","and","applies","the","EMA","params","to","the","model","during","the","test","phases",".","Saving","the","state","in","cpu","will","save","gpu","memory",",","but","will","make","training","slower","since","the","model","parameters","will","need","to","be","moved","to","cpu","before","the","averaging",".","Note",":","This","hooks","stores","two","additional","copies","of","the","model","'s","parameters",",","which","will","increase","memory","usage","significantly","."]
docstr: The constructor method of ExponentialMovingAverageModelHook.
  
  Args:
      decay: EMA decay factor, should be in [0, 1]. A decay of 0 corresponds to
          always using the latest value (no EMA) and a decay of 1 corresponds to
          not updating weights after initialization.
      consider_bn_buffers: Whether to apply EMA to batch norm buffers
      device: Device to store the model state.
docstr_tok: ["The","constructor","method","of","ExponentialMovingAverageModelHook",".","Args",":","decay",":","EMA","decay","factor",",","should","be","in","[","0",",","1","]",".","A","decay","of","0","corresponds","to","always","using","the","latest","value","(","no","EMA",")","and","a","decay","of","1","corresponds","to","not","updating","weights","after","initialization",".","consider_bn_buffers",":","Whether","to","apply","EMA","to","batch","norm","buffers","device",":","Device","to","store","the","model","state","."]
code: def __init__(
      self,
      decay: float,
      consider_bn_buffers: bool = True,
      device: str = "gpu",
  ) -> None:
      """The constructor method of ExponentialMovingAverageModelHook.
  
      Args:
          decay: EMA decay factor, should be in [0, 1]. A decay of 0 corresponds to
              always using the latest value (no EMA) and a decay of 1 corresponds to
              not updating weights after initialization.
          consider_bn_buffers: Whether to apply EMA to batch norm buffers
          device: Device to store the model state.
      """
      super().__init__()
      assert 0 <= decay <= 1, "Decay should be between 0 and 1"
      assert device in [
          "cpu",
          "gpu",
      ], "Device should be one of cpu or gpu"
      self.decay: int = decay
      self.consider_bn_buffers = consider_bn_buffers
      self.device = "cuda" if device == "gpu" else "cpu"
      self.state.model_state = {}
      self.state.ema_model_state = {}
      self.ema_model_state_list = []
      self.param_list = []
      logging.info(
          f"{self.__class__.__name__} initialized with a decay of "
          f"{decay} on device {device}"
      )
  
file_path: ClassyVision/classy_vision/hooks/exponential_moving_average_model_hook.py
imports: import itertools
  import logging
  from typing import Any, Dict, Iterable, Tuple
  import torch
  import torch.nn
  from classy_vision.generic.util import get_torch_version
  from classy_vision.hooks import register_hook
  from classy_vision.hooks.classy_hook import ClassyHook
label_desc: Use image node_id 1 to create a new ExponentialMovingAverageModelHook object from inherited base classes: ClassyHook with example: obj = ExponentialMovingAverageModelHook(decay, consider_bn_buffers, device)