class_docstr: Hook that helps to freeze some specified model parameters for certain
  number of training iterations. The parameters
  are specified in a dictionary containing {param_name: frozen_iterations}.
  Used in SwAV training.
class_docstr_tok: ["Hook","that","helps","to","freeze","some","specified","model","parameters","for","certain","number","of","training","iterations",".","The","parameters","are","specified","in","a","dictionary","containing","{","param_name",":","frozen_iterations","}",".","Used","in","SwAV","training","."]
docstr: After every backward pass and before updating the parameters,
  check if there are parameters that should stay frozen.
  Set the grad to None for those params.
docstr_tok: ["After","every","backward","pass","and","before","updating","the","parameters",",","check","if","there","are","parameters","that","should","stay","frozen",".","Set","the","grad","to","None","for","those","params","."]
code: def on_backward(self, task: "tasks.ClassyTask") -> None:
      """
      After every backward pass and before updating the parameters,
      check if there are parameters that should stay frozen.
      Set the grad to None for those params.
      """
      if len(task.config.MODEL.TEMP_FROZEN_PARAMS_ITER_MAP) == 0:
          return
      map_params_to_iters = {}
      for to_map in task.config.MODEL.TEMP_FROZEN_PARAMS_ITER_MAP:
          map_params_to_iters[to_map[0]] = to_map[1]
  
      # get the maximum iterations until which the params are frozen.
      # if the iterations are past the maximum iterations freezing any
      # param, we simply return.
      max_iterations = max(list(map_params_to_iters.values()))
      if task.iteration >= max_iterations:
          if task.iteration == max_iterations:
              logging.info(
                  f"No parameters grad removed from now on: {task.iteration}"
              )
          return
  
      num_matched_named_params = 0
      for name, p in task.model.named_parameters():
          match_param_name = self._clean_param_path(name)
          if (
              match_param_name in map_params_to_iters
          ) and task.iteration < map_params_to_iters[match_param_name]:
              num_matched_named_params += 1
              p.grad = None
  
      # TODO (Min): we need to check the exact target number.
      assert num_matched_named_params > 0, (
          f"Didn't find expected number of layers: "
          f"{num_matched_named_params} vs.  {len(map_params_to_iters)}"
      )
  
imports: import logging
  import time
  import torch
  from classy_vision import tasks
  from classy_vision.generic.profiler import compute_activations, compute_flops, count_params
  from classy_vision.hooks.classy_hook import ClassyHook
  from fvcore.common.file_io import PathManager
  from vissl.data import AirstoreDataset, GenericSSLDataset
  from vissl.models.model_helpers import model_output_has_nan
  from vissl.utils.env import get_machine_local_and_dist_rank
label_desc: Use image node_id 1 for calling the FreezeParametersHook obj's underlying member method code with example usage: obj.on_backward(task) without return types