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 calculates the scaling factor for preceding layer's grad inputs
  and current layer's grad outputs. These values are used to register a full backward
  hook. The handle returned from registering the backward hook is appended to a list
  of handles. New hooks are created and registered at every step and a new list of
  handles is created. The handles are flushed out in the unscale function.
docstr_tok: ["For","each","layer","calculates","the","scaling","factor","for","preceding","layer","'s","grad","inputs","and","current","layer","'s","grad","outputs",".","These","values","are","used","to","register","a","full","backward","hook",".","The","handle","returned","from","registering","the","backward","hook","is","appended","to","a","list","of","handles",".","New","hooks","are","created","and","registered","at","every","step","and","a","new","list","of","handles","is","created",".","The","handles","are","flushed","out","in","the","unscale","function","."]
code: def scale(self) -> None:
      """
      For each layer calculates the scaling factor for preceding layer's grad inputs
      and current layer's grad outputs. These values are used to register a full backward
      hook. The handle returned from registering the backward hook is appended to a list
      of handles. New hooks are created and registered at every step and a new list of
      handles is created. The handles are flushed out in the unscale function.
      """
      if not self._apply_layerwise_scaling:
          return
  
      for idx in range(len(self.layer_info)):
          elt = self.layer_info[idx]
          layer_name, layer = elt.layer_name, elt.layer
  
          inputs_multiplier = 1.0
          if idx > 0:
              inputs_multiplier = self.layer_info[
                  idx - 1
              ].scaling_factor
  
          outputs_multiplier = 1.0 / elt.scaling_factor
          helper = GradientHelper(
              layer_name, inputs_multiplier, outputs_multiplier
          )
          layer_handle = layer.register_full_backward_hook(
              helper.scale_gradients
          )
          self._handles.append(layer_handle)
          logging.debug(
              "name = %s \t scale = %s"
              % (layer_name, elt.scaling_factor)
          )
  
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 3 for calling the LayerwiseGradientScaler obj's underlying member method code with example usage: obj.scale() without return types