class_docstr: LayerwiseGradientScaler enables using distinct scaling factors for each layer
  of the network.
  
  Example:
  
  # Create a convolutional network
      class ConvNet(nn.Module):
          def __init__(self):
              ...
  
          def forward(self, x):
              ...
  
      # Create an instance of the model
      model = ConvNet()
      optimizer = torch.optim.SGD(model.parameters())
  
      # specify the layers to scale and their scaling factor
      layer_scale_dict = {"conv1": 2**10, "conv2": 2**8, "fc1": 2**10, "fc2": 2**9}
      scaler = LayerwiseGradientScaler(model, layer_scale_dict)
  
      for epoch in num_epochs:
          for inputs, targets in batch:
              optimizer.zero_grad()
  
              # scale the gradients
              scaler.scale()
  
              # enables mixed precision training
              with autocast():
                  predictions = model(inputs)
                  loss = loss_function(predictions, targets)
  
              loss.backward()
  
              # unscale the gradients
              loss.unscale()
  
              # step is taken if there are no inf/nan in the gradients
              # scaling factor for each layer are updated
              loss.step(optimizer)
  
  Args:
      model                      : instance of a Model class, such as ConvNet above
      layer_scale_dict (dict)    : dictionary with key = layer_name and value = scaling_factor
      growth_factor (float)      : per layer scaling factor multiplier
      backoff_factor (float)     : per layer scaling factor multiplier when an inf/nan is found
      growth_interval (int)      : number of steps after which scale is multiplied by growth_factor
      min_scaling_factor (float) : smallest scaling factor
      max_scaling_factor (float) : largest scaling factor
class_docstr_tok: ["LayerwiseGradientScaler","enables","using","distinct","scaling","factors","for","each","layer","of","the","network",".","Example",":","#","Create","a","convolutional","network","class","ConvNet","(","nn.Module",")",":","def","__init__","(","self",")",":","...","def","forward","(","self",",","x",")",":","...","#","Create","an","instance","of","the","model","model","=","ConvNet","(",")","optimizer","=","torch.optim.SGD","(","model.parameters","(",")",")","#","specify","the","layers","to","scale","and","their","scaling","factor","layer_scale_dict","=","{","``","conv1","''",":","2","*","*","10",",","``","conv2","''",":","2","*","*","8",",","``","fc1","''",":","2","*","*","10",",","``","fc2","''",":","2","*","*","9","}","scaler","=","LayerwiseGradientScaler","(","model",",","layer_scale_dict",")","for","epoch","in","num_epochs",":","for","inputs",",","targets","in","batch",":","optimizer.zero_grad","(",")","#","scale","the","gradients","scaler.scale","(",")","#","enables","mixed","precision","training","with","autocast","(",")",":","predictions","=","model","(","inputs",")","loss","=","loss_function","(","predictions",",","targets",")","loss.backward","(",")","#","unscale","the","gradients","loss.unscale","(",")","#","step","is","taken","if","there","are","no","inf\/nan","in","the","gradients","#","scaling","factor","for","each","layer","are","updated","loss.step","(","optimizer",")","Args",":","model",":","instance","of","a","Model","class",",","such","as","ConvNet","above","layer_scale_dict","(","dict",")",":","dictionary","with","key","=","layer_name","and","value","=","scaling_factor","growth_factor","(","float",")",":","per","layer","scaling","factor","multiplier","backoff_factor","(","float",")",":","per","layer","scaling","factor","multiplier","when","an","inf\/nan","is","found","growth_interval","(","int",")",":","number","of","steps","after","which","scale","is","multiplied","by","growth_factor","min_scaling_factor","(","float",")",":","smallest","scaling","factor","max_scaling_factor","(","float",")",":","largest","scaling","factor"]
docstr: For each layer, if an inf/nan is found, then multiply the scaling factor
  of that layer by the backoff factor and set the growth tracker of that
  layer to 0. Else, increment the growth tracker of the layer. If growth
  tracker equals the growth interval, then multiply the scaling factor of
  the layer by the growth factor and reset the layer's growth tracker to 0.
  Finally, clip the scaling factor to the range
  [self.min_scaling_factor, self.max_scaling_factor]. The min/max scaling
  factor values are user configurable.
docstr_tok: ["For","each","layer",",","if","an","inf\/nan","is","found",",","then","multiply","the","scaling","factor","of","that","layer","by","the","backoff","factor","and","set","the","growth","tracker","of","that","layer","to","0",".","Else",",","increment","the","growth","tracker","of","the","layer",".","If","growth","tracker","equals","the","growth","interval",",","then","multiply","the","scaling","factor","of","the","layer","by","the","growth","factor","and","reset","the","layer","'s","growth","tracker","to","0",".","Finally",",","clip","the","scaling","factor","to","the","range","[","self.min_scaling_factor",",","self.max_scaling_factor","]",".","The","min\/max","scaling","factor","values","are","user","configurable","."]
code: def _update_scale(self) -> None:
      """
      For each layer, if an inf/nan is found, then multiply the scaling factor
      of that layer by the backoff factor and set the growth tracker of that
      layer to 0. Else, increment the growth tracker of the layer. If growth
      tracker equals the growth interval, then multiply the scaling factor of
      the layer by the growth factor and reset the layer's growth tracker to 0.
      Finally, clip the scaling factor to the range
      [self.min_scaling_factor, self.max_scaling_factor]. The min/max scaling
      factor values are user configurable.
      """
      if not self._apply_layerwise_scaling:
          return
  
      for layer in self.layer_info:
          if layer.found_inf_or_nan:
              if layer.scale_layer:
                  layer.scaling_factor = max(
                      self._min_scale,
                      min(
                          self._backoff_factor * layer.scaling_factor,
                          self._max_scale,
                      ),
                  )
                  layer.growth_tracker = 0
          else:
              layer.growth_tracker += 1
              if (
                  layer.scale_layer
                  and layer.growth_tracker == self._growth_interval
              ):
                  layer.scaling_factor = max(
                      self._min_scale,
                      min(
                          self._growth_factor * layer.scaling_factor,
                          self._max_scale,
                      ),
                  )
                  layer.growth_tracker = 0
  
file_path: fairscale/fairscale/optim/layerwise_gradient_scaler.py
imports: import logging
  from typing import List, Tuple
  import torch
  import torch.nn
label_desc: Use image node_id 8 for calling the LayerwiseGradientScaler obj's underlying member method code with example usage: obj._update_scale() without return types