class_docstr: Weight maskser class for AMC pruner. Currently, AMCPruner only supports pruning kernel
  size 1x1 pointwise Conv2d layer. Before using this class to prune kernels, AMCPruner
  collected input and output feature maps for each layer, the features maps are flattened
  and save into wrapper.input_feat and wrapper.output_feat.
  
  Parameters
  ----------
  model: nn.Module
      model to be pruned
  pruner: Pruner
      A Pruner instance used to prune the model
  preserve_round: int
      after pruning, preserve filters/channels round to `preserve_round`, for example:
      for a Conv2d layer, output channel is 32, sparsity is 0.2, if preserve_round is
      1 (no preserve round), then there will be int(32 * 0.2) = 6 filters pruned, and
      32 - 6 = 26 filters are preserved. If preserve_round is 4, preserved filters will
      be round up to 28 (which can be divided by 4) and only 4 filters are pruned.
class_docstr_tok: ["Weight","maskser","class","for","AMC","pruner",".","Currently",",","AMCPruner","only","supports","pruning","kernel","size","1x1","pointwise","Conv2d","layer",".","Before","using","this","class","to","prune","kernels",",","AMCPruner","collected","input","and","output","feature","maps","for","each","layer",",","the","features","maps","are","flattened","and","save","into","wrapper.input_feat","and","wrapper.output_feat",".","Parameters","--","--","--","--","--","model",":","nn.Module","model","to","be","pruned","pruner",":","Pruner","A","Pruner","instance","used","to","prune","the","model","preserve_round",":","int","after","pruning",",","preserve","filters\/channels","round","to","`","preserve_round","`",",","for","example",":","for","a","Conv2d","layer",",","output","channel","is","32",",","sparsity","is","0.2",",","if","preserve_round","is","1","(","no","preserve","round",")",",","then","there","will","be","int","(","32","*","0.2",")","=","6","filters","pruned",",","and","32","-","6","=","26","filters","are","preserved",".","If","preserve_round","is","4",",","preserved","filters","will","be","round","up","to","28","(","which","can","be","divided","by","4",")","and","only","4","filters","are","pruned","."]
docstr: Calculate the mask of given layer.
  Parameters
  ----------
  sparsity: float
      pruning ratio,  preserved weight ratio is `1 - sparsity`
  wrapper: PrunerModuleWrapper
      layer wrapper of this layer
  wrapper_idx: int
      index of this wrapper in pruner's all wrappers
  Returns
  -------
  dict
      dictionary for storing masks, keys of the dict:
      'weight_mask':  weight mask tensor
      'bias_mask': bias mask tensor (optional)
docstr_tok: ["Calculate","the","mask","of","given","layer",".","Parameters","--","--","--","--","--","sparsity",":","float","pruning","ratio",",","preserved","weight","ratio","is","`","1","-","sparsity","`","wrapper",":","PrunerModuleWrapper","layer","wrapper","of","this","layer","wrapper_idx",":","int","index","of","this","wrapper","in","pruner","'s","all","wrappers","Returns","--","--","--","-","dict","dictionary","for","storing","masks",",","keys","of","the","dict",":","'weight_mask","'",":","weight","mask","tensor","'bias_mask","'",":","bias","mask","tensor","(","optional",")"]
code: def calc_mask(
      self, sparsity, wrapper, wrapper_idx=None, preserve_idx=None
  ):
      """
      Calculate the mask of given layer.
      Parameters
      ----------
      sparsity: float
          pruning ratio,  preserved weight ratio is `1 - sparsity`
      wrapper: PrunerModuleWrapper
          layer wrapper of this layer
      wrapper_idx: int
          index of this wrapper in pruner's all wrappers
      Returns
      -------
      dict
          dictionary for storing masks, keys of the dict:
          'weight_mask':  weight mask tensor
          'bias_mask': bias mask tensor (optional)
      """
      msg = "module type {} is not supported!".format(wrapper.type)
      assert wrapper.type in ["Conv2d", "Linear"], msg
      weight = wrapper.module.weight.data
      bias = None
      if (
          hasattr(wrapper.module, "bias")
          and wrapper.module.bias is not None
      ):
          bias = wrapper.module.bias.data
  
      if wrapper.weight_mask is None:
          mask_weight = (
              torch.ones(weight.size()).type_as(weight).detach()
          )
      else:
          mask_weight = wrapper.weight_mask.clone()
      if bias is not None:
          if wrapper.bias_mask is None:
              mask_bias = torch.ones(bias.size()).type_as(bias).detach()
          else:
              mask_bias = wrapper.bias_mask.clone()
      else:
          mask_bias = None
      mask = {"weight_mask": mask_weight, "bias_mask": mask_bias}
  
      num_total = weight.size(1)
      num_prune = int(num_total * sparsity)
      if self.preserve_round > 1:
          num_preserve = num_total - num_prune
          num_preserve = int(
              math.ceil(num_preserve * 1.0 / self.preserve_round)
              * self.preserve_round
          )
          if num_preserve > num_total:
              num_preserve = num_total
          num_prune = num_total - num_preserve
  
      if (num_total < 2 or num_prune < 1) and preserve_idx is None:
          return mask
  
      return self.get_mask(
          mask, weight, num_preserve, wrapper, wrapper_idx, preserve_idx
      )
  
imports: import logging
  import math
  import numpy
  import torch
  from .weight_masker import WeightMasker
label_desc: Use image node_id 2 for calling the AMCWeightMasker obj's underlying member method code with example usage: obj.calc_mask(sparsity, wrapper, wrapper_idx, preserve_idx) and returns: self, mask