docstr: Call this method to modify your model to take advantage of sparse matrix multiplication.
  Note that this call alone only augments the model with additional buffers needed for sparse MMA,
  it does not enable use of sparse MMA. 
  
  If you are starting with a fresh model:
  
  model = ...
  ASP.init_model_for_pruning(model, mask_calculator, ...)
  if (training) ASP.init_optimizer_for_pruning(optimizer)
  ASP.compute_sparse_masks() // sparsity is off by default, call when youy want to enable it.
  
  If you are starting from a checkpoint:
  
  model = ...
  ASP.init_model_for_pruning(model, mask_calculator, ...)
  torch.load(...)
  if (training) ASP.init_optimizer_for_pruning(optimizer)
  
  Arguments:
    model                    The model
    mask_calculator          Either callable that computes mask given a tensor OR pattern string for sparse mask lib.
    verbosity                Integer controling verbosity level.
                             0 -> Only errors.
                             1 -> Errors and warnings.
                             2 -> Errors, warnings and info.
                             3 -> Errors, warnings, info and debug.
    whitelist                Module types approved for sparsity.
    allowed_layer_names      If not None, only layer names that appear in this list are considered for sparsity.
    disallowed_layer_names   If not [], only layer names that do not appear in this list are considered for sparsity.
    allow_recompute_mask     If True, stores pruned values so that dense weights can be restored.
                             Pruned weights are stored in CPU memory, hence this option does not increase GPU memory usage.
    custom_layer_dict        Dictionary of additional layer paremeters to sparsify. e.g. {CustomLinear: ['weight']}
    allow_permutation        If True, allow the input channel permutation to ease the influence of weight pruning.
    
    [Future] Support for allow_recompute_mask can be removed, it is not part of sparse inference recipe.
docstr_tok: ["Call","this","method","to","modify","your","model","to","take","advantage","of","sparse","matrix","multiplication",".","Note","that","this","call","alone","only","augments","the","model","with","additional","buffers","needed","for","sparse","MMA",",","it","does","not","enable","use","of","sparse","MMA",".","If","you","are","starting","with","a","fresh","model",":","model","=","...","ASP.init_model_for_pruning","(","model",",","mask_calculator",",","...",")","if","(","training",")","ASP.init_optimizer_for_pruning","(","optimizer",")","ASP.compute_sparse_masks","(",")","\/\/","sparsity","is","off","by","default",",","call","when","youy","want","to","enable","it",".","If","you","are","starting","from","a","checkpoint",":","model","=","...","ASP.init_model_for_pruning","(","model",",","mask_calculator",",","...",")","torch.load","(","...",")","if","(","training",")","ASP.init_optimizer_for_pruning","(","optimizer",")","Arguments",":","model","The","model","mask_calculator","Either","callable","that","computes","mask","given","a","tensor","OR","pattern","string","for","sparse","mask","lib",".","verbosity","Integer","controling","verbosity","level",".","0","-",">","Only","errors",".","1","-",">","Errors","and","warnings",".","2","-",">","Errors",",","warnings","and","info",".","3","-",">","Errors",",","warnings",",","info","and","debug",".","whitelist","Module","types","approved","for","sparsity",".","allowed_layer_names","If","not","None",",","only","layer","names","that","appear","in","this","list","are","considered","for","sparsity",".","disallowed_layer_names","If","not","[","]",",","only","layer","names","that","do","not","appear","in","this","list","are","considered","for","sparsity",".","allow_recompute_mask","If","True",",","stores","pruned","values","so","that","dense","weights","can","be","restored",".","Pruned","weights","are","stored","in","CPU","memory",",","hence","this","option","does","not","increase","GPU","memory","usage",".","custom_layer_dict","Dictionary","of","additional","layer","paremeters","to","sparsify",".","e.g",".","{","CustomLinear",":","[","'weight","'","]","}","allow_permutation","If","True",",","allow","the","input","channel","permutation","to","ease","the","influence","of","weight","pruning",".","[","Future","]","Support","for","allow_recompute_mask","can","be","removed",",","it","is","not","part","of","sparse","inference","recipe","."]
code: def init_model_for_pruning(
      cls,
      model,
      mask_calculator="m4n2_1d",
      verbosity=3,
      whitelist=[
          torch.nn.Linear,
          torch.nn.Conv1d,
          torch.nn.Conv2d,
          torch.nn.Conv3d,
          torch.nn.MultiheadAttention,
      ],
      allowed_layer_names=None,
      disallowed_layer_names=[],
      allow_recompute_mask=False,
      custom_layer_dict={},
      allow_permutation=True,
  ):
      """Call this method to modify your model to take advantage of sparse matrix multiplication.
      Note that this call alone only augments the model with additional buffers needed for sparse MMA,
      it does not enable use of sparse MMA.
  
      If you are starting with a fresh model:
  
      model = ...
      ASP.init_model_for_pruning(model, mask_calculator, ...)
      if (training) ASP.init_optimizer_for_pruning(optimizer)
      ASP.compute_sparse_masks() // sparsity is off by default, call when youy want to enable it.
  
      If you are starting from a checkpoint:
  
      model = ...
      ASP.init_model_for_pruning(model, mask_calculator, ...)
      torch.load(...)
      if (training) ASP.init_optimizer_for_pruning(optimizer)
  
      Arguments:
        model                    The model
        mask_calculator          Either callable that computes mask given a tensor OR pattern string for sparse mask lib.
        verbosity                Integer controling verbosity level.
                                 0 -> Only errors.
                                 1 -> Errors and warnings.
                                 2 -> Errors, warnings and info.
                                 3 -> Errors, warnings, info and debug.
        whitelist                Module types approved for sparsity.
        allowed_layer_names      If not None, only layer names that appear in this list are considered for sparsity.
        disallowed_layer_names   If not [], only layer names that do not appear in this list are considered for sparsity.
        allow_recompute_mask     If True, stores pruned values so that dense weights can be restored.
                                 Pruned weights are stored in CPU memory, hence this option does not increase GPU memory usage.
        custom_layer_dict        Dictionary of additional layer paremeters to sparsify. e.g. {CustomLinear: ['weight']}
        allow_permutation        If True, allow the input channel permutation to ease the influence of weight pruning.
  
        [Future] Support for allow_recompute_mask can be removed, it is not part of sparse inference recipe.
      """
      assert cls.__model is None, "ASP has been initialized already."
      cls.__model = model
      cls.__verbosity = verbosity
      cls.__allow_permutation = allow_permutation
  
      if isinstance(mask_calculator, str):
  
          def create_mask_from_pattern(param):
              return create_mask(param, mask_calculator).bool()
  
          cls.__calculate_mask = create_mask_from_pattern
      else:
          cls.__calculate_mask = (
              mask_calculator  # user defined function
          )
  
      # function to extract variables that will be sparsified.
      # idea is that you will add one of these functions for each module type that can be sparsified.
      if torchvision_imported:
          print(
              "[ASP] torchvision is imported, can work with the MaskRCNN/KeypointRCNN from torchvision."
          )
          torchvision_version = str(torchvision.__version__)
          torchvision_version_major = int(
              torchvision_version.split(".")[0]
          )
          torchvision_version_minor = int(
              torchvision_version.split(".")[1]
          )
          if (
              torchvision_version_major == 0
              and torchvision_version_minor < 12
          ):
              sparse_parameter_list = {
                  torch.nn.Linear: ["weight"],
                  torch.nn.Conv1d: ["weight"],
                  torch.nn.Conv2d: ["weight"],
                  torch.nn.Conv3d: ["weight"],
                  torch.nn.modules.linear.NonDynamicallyQuantizableLinear: [
                      "weight"
                  ],
                  torch.nn.MultiheadAttention: [
                      "q_proj_weight",
                      "k_proj_weight",
                      "v_proj_weight",
                      "in_proj_weight",
                  ],
                  torchvision.ops.misc.Conv2d: ["weight"],
              }
          else:  # Torchvision remove APIs that were deprecated before 0.8 (#5386) in 0.12.0, torchvision.ops.misc.Conv2d is removed
              sparse_parameter_list = {
                  torch.nn.Linear: ["weight"],
                  torch.nn.Conv1d: ["weight"],
                  torch.nn.Conv2d: ["weight"],
                  torch.nn.Conv3d: ["weight"],
                  torch.nn.modules.linear.NonDynamicallyQuantizableLinear: [
                      "weight"
                  ],
                  torch.nn.MultiheadAttention: [
                      "q_proj_weight",
                      "k_proj_weight",
                      "v_proj_weight",
                      "in_proj_weight",
                  ],
              }
      else:
          sparse_parameter_list = {
              torch.nn.Linear: ["weight"],
              torch.nn.Conv1d: ["weight"],
              torch.nn.Conv2d: ["weight"],
              torch.nn.Conv3d: ["weight"],
              torch.nn.modules.linear.NonDynamicallyQuantizableLinear: [
                  "weight"
              ],
              torch.nn.MultiheadAttention: [
                  "q_proj_weight",
                  "k_proj_weight",
                  "v_proj_weight",
                  "in_proj_weight",
              ],
          }
      if (
          custom_layer_dict
      ):  # Update default list to include user supplied custom (layer type : parameter tensor), make sure this tensor type is something ASP knows how to prune
          sparse_parameter_list.update(custom_layer_dict)
          whitelist += list(custom_layer_dict.keys())
  
      for module_type in whitelist:
          assert module_type in sparse_parameter_list, (
              "Module %s :: Don't know how to sparsify module."
              % module.dtype()
          )
  
      # find all sparse modules, extract sparse parameters and decorate
      def add_sparse_attributes(module_name, module):
          sparse_parameters = sparse_parameter_list[type(module)]
          for p_name, p in module.named_parameters():
              if p_name in sparse_parameters and p.requires_grad:
                  # check for NVIDIA's TC compatibility: we check along the horizontal direction
                  if p.dtype == torch.float32 and (
                      (p.size()[0] % 8) != 0 or (p.size()[1] % 16) != 0
                  ):  # User defines FP32 and APEX internally uses FP16 math
                      print(
                          "[ASP] Auto skipping pruning %s::%s of size=%s and type=%s for sparsity"
                          % (
                              module_name,
                              p_name,
                              str(p.size()),
                              str(p.dtype),
                          )
                      )
                      continue
                  if p.dtype == torch.float16 and (
                      (p.size()[0] % 8) != 0 or (p.size()[1] % 16) != 0
                  ):  # For Conv2d dim= K x CRS; we prune along C
                      print(
                          "[ASP] Auto skipping pruning %s::%s of size=%s and type=%s for sparsity"
                          % (
                              module_name,
                              p_name,
                              str(p.size()),
                              str(p.dtype),
                          )
                      )
                      continue
  
                  if cls.__verbosity >= 3:
                      print(
                          "[ASP] Sparsifying %s::%s of size=%s and type=%s for sparsity"
                          % (
                              module_name,
                              p_name,
                              str(p.size()),
                              str(p.dtype),
                          )
                      )
  
                  mask = torch.ones_like(p).bool()
                  buffname = p_name.split(".")[
                      -1
                  ]  # buffer names cannot contain "."
                  module.register_buffer(
                      "__%s_mma_mask" % buffname, mask
                  )
                  if allow_recompute_mask:
                      pruned = torch.zeros_like(p).cpu()
                      module.register_buffer(
                          "__%s_mma_pruned_p" % buffname, pruned
                      )
                  else:
                      pruned = None
                  cls.__sparse_parameters.append(
                      (module_name, module, p_name, p, mask, pruned)
                  )
              else:
                  if cls.__verbosity >= 3:
                      print(
                          "[ASP] Not sparsifying %s::%s of size=%s and type=%s"
                          % (
                              module_name,
                              p_name,
                              str(p.size()),
                              str(p.dtype),
                          )
                      )
  
      for name, sparse_module in eligible_modules(
          model,
          tuple(whitelist),
          allowed_layer_names,
          disallowed_layer_names,
      ):
          add_sparse_attributes(name, sparse_module)
  
      if (
          allow_permutation
      ):  # find all named modules, extract parameters and decorate, used for offline permutation in K dim
          for module_name, module in model.named_modules():
              module_type_str = str(type(module)).split("'")[1]
              if (
                  module_type_str
                  == "torch.nn.modules.container.Sequential"
                  or module_type_str.startswith("torchvision.models")
              ):
                  # filter out the 'torch.nn.modules.container.Sequential' type and the whole model, like 'torchvision.models.vgg.VGG'
                  continue
              for p_name, p in module.named_parameters():
                  cls.__all_parameters.append(
                      (module_name, module, p_name, p)
                  )
              if (
                  module_type_str
                  == "torch.nn.modules.batchnorm.BatchNorm2d"
              ):
                  # need to get the running_mean and running_var from model.state_dict(), as they are not the learnable parameters
                  module_mean_name = module_name + ".running_mean"
                  module_var_name = module_name + ".running_var"
                  for param_key in model.state_dict():
                      if (
                          module_mean_name == param_key
                          or module_var_name == param_key
                      ):
                          cls.__all_parameters.append(
                              (
                                  module_name,
                                  module,
                                  param_key.split(".")[-1],
                                  model.state_dict()[param_key],
                              )
                          )
          # add the __permutation_output_dir field to save the intermediate results for permutation
          cls.__permutation_output_dir = "."
          # Set the corresponding params from ASP class to the Permutation class
          permutation_verbosity = 5
          Permutation.set_permutation_params_from_asp(
              cls.__model,
              cls.__sparse_parameters,
              cls.__all_parameters,
              permutation_verbosity,
          )
          # Set the identical random seed for all GPUs to make sure the same results generated in permutation search
          Permutation.set_identical_seed()
  
file_path: apex/apex/contrib/sparsity/asp.py
imports: import types
  import torch
  from .sparse_masklib import create_mask
  from .permutation_lib import Permutation
  import json
  import os
  import string
  import time
label_desc: Use image node_id 1 for calling the ASP obj's underlying member method code with example usage: obj.init_model_for_pruning(cls, model, mask_calculator, verbosity, whitelist, allowed_layer_names, disallowed_layer_names, allow_recompute_mask, custom_layer_dict, allow_permutation) and returns: create_mask