docstr: Call this method to enable sparsity.
  If init(...) was called with allow_recompute_mask=False AND sparsity is disabled, pruned field can be None.
docstr_tok: ["Call","this","method","to","enable","sparsity",".","If","init","(","...",")","was","called","with","allow_recompute_mask=False","AND","sparsity","is","disabled",",","pruned","field","can","be","None","."]
code: def compute_sparse_masks(cls):
      """Call this method to enable sparsity.
      If init(...) was called with allow_recompute_mask=False AND sparsity is disabled, pruned field can be None.
      """
      with torch.no_grad():
          if cls.__allow_permutation:
              # Step 1: use the Torch.FX library to build the graph
              # Step 2: permutation search with the customized kernel
              # The simplest without user intervention:
              # A. try to import with the distributed mode of the original model
              # B. if meet the error, import with the none-distributed mode of the original model
              start_time_permute = time.perf_counter()
              successful_permutation = False
              try:
                  successful_permutation = Permutation.permute_model(
                      cls.__model.module,
                      dump_fx_graph=cls.__save_permutation_graph,
                      save_dumped_fx_graph=os.path.join(
                          cls.__permutation_output_dir,
                          "model_offline_permutation_graph.json",
                      ),
                  )
                  if successful_permutation:
                      print(
                          "\n[compute_sparse_masks] permuted the (distributed) model."
                      )
              except AttributeError:
                  successful_permutation = Permutation.permute_model(
                      cls.__model,
                      dump_fx_graph=cls.__save_permutation_graph,
                      save_dumped_fx_graph=os.path.join(
                          cls.__permutation_output_dir,
                          "model_offline_permutation_graph.json",
                      ),
                  )
                  if successful_permutation:
                      print(
                          "\n[compute_sparse_masks] permuted the model."
                      )
  
              if successful_permutation:
                  duration_build_offline_permutation_graph = (
                      time.perf_counter() - start_time_permute
                  )
                  print(
                      "[compute_sparse_masks] Take {:.4f} seconds to find and apply permutations.".format(
                          duration_build_offline_permutation_graph
                      )
                  )
  
          for (
              module_name,
              module,
              p_name,
              p,
              mask,
              pruned,
          ) in cls.__sparse_parameters:
              if mask.sum() < mask.numel():  # when recalculating masks
                  # restore dense parameter if allow_recompute_mask is enabled
                  assert (
                      pruned is not None
                  ), "Unable to restore dense parameter because allow_recompute_mask == False"
                  p.add_(pruned.cuda())
  
              mask.set_(cls.__calculate_mask(p))
  
              if pruned is not None:  # stow away pruned weights to cpu
                  pruned.set_((p * (~mask)).cpu())
  
              p.mul_(
                  mask
              )  # in-place multiplication, so pruned weights are 0-values, hence checkpoint will have 0s for pruned weights
              if cls.__verbosity >= 2:
                  print(
                      "[ASP] Enabled %.2f%% sparsity for %s::%s of size=%s and type=%s with magnitude %s"
                      % (
                          100.0 - 100.0 * mask.sum() / mask.numel(),
                          module_name,
                          p_name,
                          str(p.size()),
                          str(p.dtype),
                          torch.sum(torch.abs(p)),
                      )
                  )
  
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 4 for calling the ASP obj's underlying member method code with example usage: obj.compute_sparse_masks(cls) without return types