class_docstr: "
  This class implement the second order optimizer - KFAC, which uses Kronecker factor products of inputs and the gradients to
  get the approximate inverse fisher matrix, which is used to update the model parameters. Presently this optimizer works only
  on liner and 2D convolution layers. If you want to know more details about KFAC, please check the paper [1]_ and [2]_.
  
  References:
  -----------
  [1] Martens, James, and Roger Grosse. Optimizing Neural Networks with Kronecker-Factored Approximate Curvature.
  arXiv:1503.05671, arXiv, 7 June 2020. arXiv.org, http://arxiv.org/abs/1503.05671.
  [2] Grosse, Roger, and James Martens. A Kronecker-Factored Approximate Fisher Matrix for Convolution Layers.
  arXiv:1602.01407, arXiv, 23 May 2016. arXiv.org, http://arxiv.org/abs/1602.01407.
class_docstr_tok: ["``","This","class","implement","the","second","order","optimizer","-","KFAC",",","which","uses","Kronecker","factor","products","of","inputs","and","the","gradients","to","get","the","approximate","inverse","fisher","matrix",",","which","is","used","to","update","the","model","parameters",".","Presently","this","optimizer","works","only","on","liner","and","2D","convolution","layers",".","If","you","want","to","know","more","details","about","KFAC",",","please","check","the","paper","[","1","]","_","and","[","2","]","_.","References",":","--","--","--","--","--","-","[","1","]","Martens",",","James",",","and","Roger","Grosse",".","Optimizing","Neural","Networks","with","Kronecker-Factored","Approximate","Curvature",".","arXiv:1503.05671",",","arXiv",",","7","June","2020.","arXiv.org",",","http",":","\/\/arxiv.org\/abs\/1503.05671",".","[","2","]","Grosse",",","Roger",",","and","James","Martens",".","A","Kronecker-Factored","Approximate","Fisher","Matrix","for","Convolution","Layers",".","arXiv:1602.01407",",","arXiv",",","23","May","2016.","arXiv.org",",","http",":","\/\/arxiv.org\/abs\/1602.01407","."]
docstr: Parameters:
  -----------
  model: torch.nn.Module
      The model to be optimized.
  lr: float (default: 0.001)
      Learning rate for the optimizer.
  momentum: float (default: 0.9)
      Momentum for the optimizer.
  stat_decay: float (default: 0.95)
      Decay rate for the update of covariance matrix with mean.
  damping: float (default: 0.001)
      damping factor for the update of covariance matrix.
  kl_clip: float (default: 0.001)
      Clipping value for the update of covariance matrix.
  weight_decay: float (default: 0)
      weight decay for the optimizer.
  Tcov: int (default: 10)
      The number of steps to update the covariance matrix.
  Tinv: int (default: 100)
      The number of steps to calculate the inverse of covariance matrix.
  batch_averaged: bool (default: True)
      States whether to use batch averaged covariance matrix.
  mean: bool (default: False)
      States whether to use mean centered covariance matrix.
docstr_tok: ["Parameters",":","--","--","--","--","--","-","model",":","torch.nn.Module","The","model","to","be","optimized",".","lr",":","float","(","default",":","0.001",")","Learning","rate","for","the","optimizer",".","momentum",":","float","(","default",":","0.9",")","Momentum","for","the","optimizer",".","stat_decay",":","float","(","default",":","0.95",")","Decay","rate","for","the","update","of","covariance","matrix","with","mean",".","damping",":","float","(","default",":","0.001",")","damping","factor","for","the","update","of","covariance","matrix",".","kl_clip",":","float","(","default",":","0.001",")","Clipping","value","for","the","update","of","covariance","matrix",".","weight_decay",":","float","(","default",":","0",")","weight","decay","for","the","optimizer",".","Tcov",":","int","(","default",":","10",")","The","number","of","steps","to","update","the","covariance","matrix",".","Tinv",":","int","(","default",":","100",")","The","number","of","steps","to","calculate","the","inverse","of","covariance","matrix",".","batch_averaged",":","bool","(","default",":","True",")","States","whether","to","use","batch","averaged","covariance","matrix",".","mean",":","bool","(","default",":","False",")","States","whether","to","use","mean","centered","covariance","matrix","."]
code: def __init__(
      self,
      model: torch.nn.Module,
      lr: float = 0.001,
      momentum: float = 0.9,
      stat_decay: float = 0.95,
      damping: float = 0.001,
      kl_clip: float = 0.001,
      weight_decay: float = 0,
      TCov: int = 10,
      TInv: int = 100,
      batch_averaged: bool = True,
      mean: bool = False,
  ):
      """
      Parameters:
      -----------
      model: torch.nn.Module
          The model to be optimized.
      lr: float (default: 0.001)
          Learning rate for the optimizer.
      momentum: float (default: 0.9)
          Momentum for the optimizer.
      stat_decay: float (default: 0.95)
          Decay rate for the update of covariance matrix with mean.
      damping: float (default: 0.001)
          damping factor for the update of covariance matrix.
      kl_clip: float (default: 0.001)
          Clipping value for the update of covariance matrix.
      weight_decay: float (default: 0)
          weight decay for the optimizer.
      Tcov: int (default: 10)
          The number of steps to update the covariance matrix.
      Tinv: int (default: 100)
          The number of steps to calculate the inverse of covariance matrix.
      batch_averaged: bool (default: True)
          States whether to use batch averaged covariance matrix.
      mean: bool (default: False)
          States whether to use mean centered covariance matrix.
      """
  
      if lr < 0.0:
          raise ValueError("Invalid learning rate: {}".format(lr))
      if momentum < 0.0:
          raise ValueError(
              "Invalid momentum value: {}".format(momentum)
          )
      if weight_decay < 0.0:
          raise ValueError(
              "Invalid weight_decay value: {}".format(weight_decay)
          )
      defaults = dict(
          lr=lr,
          momentum=momentum,
          damping=damping,
          weight_decay=weight_decay,
      )
      super(KFACOptimizer, self).__init__(model.parameters(), defaults)
      self.batch_averaged = batch_averaged
  
      self.known_modules = {"Linear", "Conv2d"}
  
      self.modules: List[torch.nn.Module] = []
  
      self.model = model
      self._prepare_model()
  
      self.steps = 0
  
      self.m_aa: Dict[torch.nn.Module, torch.Tensor] = {}
      self.m_gg: Dict[torch.nn.Module, torch.Tensor] = {}
      self.Q_a: Dict[torch.nn.Module, torch.Tensor] = {}
      self.Q_g: Dict[torch.nn.Module, torch.Tensor] = {}
      self.d_a: Dict[torch.nn.Module, torch.Tensor] = {}
      self.d_g: Dict[torch.nn.Module, torch.Tensor] = {}
      self.stat_decay = stat_decay
  
      self.kl_clip = kl_clip
      self.TCov = TCov
      self.TInv = TInv
  
      self.mean = mean
  
file_path: deepchem/deepchem/models/torch_models/kfac_optimizer.py
imports: import math
  from typing import Optional, Callable, Dict, Tuple, List, Union
label_desc: Use image node_id 1 to create a new KFACOptimizer object from inherited base classes: optim with example: obj = KFACOptimizer(model, lr, momentum, stat_decay, damping, kl_clip, weight_decay, TCov, TInv, batch_averaged, mean)