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, check if any of the layer's parameters contain an inf/nan.
  If there are no inf/nan in the gradient, then gradient of that layer is
  unscaled by the reciprocal of the scaling factor for that layer.
  Finally, all handles recorded while registering the hooks are deleted.
docstr_tok: ["For","each","layer",",","check","if","any","of","the","layer","'s","parameters","contain","an","inf\/nan",".","If","there","are","no","inf\/nan","in","the","gradient",",","then","gradient","of","that","layer","is","unscaled","by","the","reciprocal","of","the","scaling","factor","for","that","layer",".","Finally",",","all","handles","recorded","while","registering","the","hooks","are","deleted","."]
code: def unscale(self) -> None:
      """
      For each layer, check if any of the layer's parameters contain an inf/nan.
      If there are no inf/nan in the gradient, then gradient of that layer is
      unscaled by the reciprocal of the scaling factor for that layer.
      Finally, all handles recorded while registering the hooks are deleted.
      """
      if not self._apply_layerwise_scaling:
          return
  
      layers_with_finite_values = self._get_layers_with_finite_values()
      for item in layers_with_finite_values:
          for param_name, param in item.layer.named_parameters():
              if hasattr(param, "grad") and param.grad is not None:
                  logging.debug(
                      "%s scaling down %s by %s"
                      % (
                          item.layer_name,
                          param_name,
                          1.0 / item.scaling_factor,
                      )
                  )
                  param.grad.mul_(1.0 / item.scaling_factor)
  
      while len(self._handles) > 0:
          elt = self._handles.pop()
          elt.remove()
  
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 5 for calling the LayerwiseGradientScaler obj's underlying member method code with example usage: obj.unscale() without return types