class_docstr: A Pytorch implementation of NetAdapt compression algorithm.
  
  Parameters
  ----------
  model : pytorch model
      The model to be pruned.
  config_list : list
      Supported keys:
          - sparsity : The target overall sparsity.
          - op_types : The operation type to prune.
  short_term_fine_tuner : function
      function to short-term fine tune the masked model.
      This function should include `model` as the only parameter,
      and fine tune the model for a short term after each pruning iteration.
      Example::
  
          def short_term_fine_tuner(model, epoch=3):
              device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
              train_loader = ...
              criterion = torch.nn.CrossEntropyLoss()
              optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
              model.train()
              for _ in range(epoch):
                  for batch_idx, (data, target) in enumerate(train_loader):
                      data, target = data.to(device), target.to(device)
                      optimizer.zero_grad()
                      output = model(data)
                      loss = criterion(output, target)
                      loss.backward()
                      optimizer.step()
  evaluator : function
      function to evaluate the masked model.
      This function should include `model` as the only parameter, and returns a scalar value.
      Example::
  
          def evaluator(model):
              device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
              val_loader = ...
              model.eval()
              correct = 0
              with torch.no_grad():
                  for data, target in val_loader:
                      data, target = data.to(device), target.to(device)
                      output = model(data)
                      # get the index of the max log-probability
                      pred = output.argmax(dim=1, keepdim=True)
                      correct += pred.eq(target.view_as(pred)).sum().item()
              accuracy = correct / len(val_loader.dataset)
              return accuracy
  optimize_mode : str
      optimize mode, `maximize` or `minimize`, by default `maximize`.
  base_algo : str
      Base pruning algorithm. `level`, `l1`, `l2` or `fpgm`, by default `l1`. Given the sparsity distribution among the ops,
      the assigned `base_algo` is used to decide which filters/channels/weights to prune.
  sparsity_per_iteration : float
      sparsity to prune in each iteration.
  experiment_data_dir : str
      PATH to save experiment data,
      including the config_list generated for the base pruning algorithm and the performance of the pruned model.
class_docstr_tok: ["A","Pytorch","implementation","of","NetAdapt","compression","algorithm",".","Parameters","--","--","--","--","--","model",":","pytorch","model","The","model","to","be","pruned",".","config_list",":","list","Supported","keys",":","-","sparsity",":","The","target","overall","sparsity",".","-","op_types",":","The","operation","type","to","prune",".","short_term_fine_tuner",":","function","function","to","short-term","fine","tune","the","masked","model",".","This","function","should","include","`","model","`","as","the","only","parameter",",","and","fine","tune","the","model","for","a","short","term","after","each","pruning","iteration",".","Example",":",":","def","short_term_fine_tuner","(","model",",","epoch=3",")",":","device","=","torch.device","(","``","cuda","''","if","torch.cuda.is_available","(",")","else","``","cpu","''",")","train_loader","=","...","criterion","=","torch.nn.CrossEntropyLoss","(",")","optimizer","=","torch.optim.SGD","(","model.parameters","(",")",",","lr=0.01",")","model.train","(",")","for","_","in","range","(","epoch",")",":","for","batch_idx",",","(","data",",","target",")","in","enumerate","(","train_loader",")",":","data",",","target","=","data.to","(","device",")",",","target.to","(","device",")","optimizer.zero_grad","(",")","output","=","model","(","data",")","loss","=","criterion","(","output",",","target",")","loss.backward","(",")","optimizer.step","(",")","evaluator",":","function","function","to","evaluate","the","masked","model",".","This","function","should","include","`","model","`","as","the","only","parameter",",","and","returns","a","scalar","value",".","Example",":",":","def","evaluator","(","model",")",":","device","=","torch.device","(","``","cuda","''","if","torch.cuda.is_available","(",")","else","``","cpu","''",")","val_loader","=","...","model.eval","(",")","correct","=","0","with","torch.no_grad","(",")",":","for","data",",","target","in","val_loader",":","data",",","target","=","data.to","(","device",")",",","target.to","(","device",")","output","=","model","(","data",")","#","get","the","index","of","the","max","log-probability","pred","=","output.argmax","(","dim=1",",","keepdim=True",")","correct","+=","pred.eq","(","target.view_as","(","pred",")",")",".sum","(",")",".item","(",")","accuracy","=","correct","\/","len","(","val_loader.dataset",")","return","accuracy","optimize_mode",":","str","optimize","mode",",","`","maximize","`","or","`","minimize","`",",","by","default","`","maximize","`",".","base_algo",":","str","Base","pruning","algorithm",".","`","level","`",",","`","l1","`",",","`","l2","`","or","`","fpgm","`",",","by","default","`","l1","`",".","Given","the","sparsity","distribution","among","the","ops",",","the","assigned","`","base_algo","`","is","used","to","decide","which","filters\/channels\/weights","to","prune",".","sparsity_per_iteration",":","float","sparsity","to","prune","in","each","iteration",".","experiment_data_dir",":","str","PATH","to","save","experiment","data",",","including","the","config_list","generated","for","the","base","pruning","algorithm","and","the","performance","of","the","pruned","model","."]
docstr: Compress the model.
  
  Returns
  -------
  torch.nn.Module
      model with specified modules compressed.
docstr_tok: ["Compress","the","model",".","Returns","--","--","--","-","torch.nn.Module","model","with","specified","modules","compressed","."]
code: def compress(self):
      """
      Compress the model.
  
      Returns
      -------
      torch.nn.Module
          model with specified modules compressed.
      """
      _logger.info("Starting NetAdapt Compression...")
  
      pruning_iteration = 0
      current_sparsity = 0
      delta_num_weights_per_iteration = int(
          get_total_num_weights(
              self._model_to_prune, ["Conv2d", "Linear"]
          )
          * self._sparsity_per_iteration
      )
  
      # stop condition
      while current_sparsity < self._sparsity:
          _logger.info("Pruning iteration: %d", pruning_iteration)
  
          # calculate target sparsity of this iteration
          target_sparsity = (
              current_sparsity + self._sparsity_per_iteration
          )
  
          # variable to store the info of the best layer found in this iteration
          best_op = {}
  
          for wrapper in self.get_modules_wrapper():
              _logger.debug("op name : %s", wrapper.name)
              _logger.debug(
                  "op weights : %d", wrapper.weight_mask.numel()
              )
              _logger.debug(
                  "op left weights : %d",
                  wrapper.weight_mask.sum().item(),
              )
  
              current_op_sparsity = (
                  1
                  - wrapper.weight_mask.sum().item()
                  / wrapper.weight_mask.numel()
              )
              _logger.debug(
                  "current op sparsity : %s", current_op_sparsity
              )
  
              # sparsity that this layer needs to prune to satisfy the requirement
              target_op_sparsity = (
                  current_op_sparsity
                  + delta_num_weights_per_iteration
                  / self._calc_num_related_weights(wrapper.name)
              )
  
              if target_op_sparsity >= 1:
                  _logger.info(
                      "Layer %s has no enough weights (remained) to prune",
                      wrapper.name,
                  )
                  continue
  
              config_list = self._update_config_list(
                  self._config_list_generated,
                  wrapper.name,
                  target_op_sparsity,
              )
              _logger.debug("config_list used : %s", config_list)
  
              pruner = PRUNER_DICT[self._base_algo](
                  copy.deepcopy(self._model_to_prune), config_list
              )
              model_masked = pruner.compress()
  
              # Short-term fine tune the pruned model
              self._short_term_fine_tuner(model_masked)
  
              performance = self._evaluator(model_masked)
              _logger.info(
                  "Layer : %s, evaluation result after short-term fine tuning : %s",
                  wrapper.name,
                  performance,
              )
  
              if (
                  not best_op
                  or (
                      self._optimize_mode is OptimizeMode.Maximize
                      and performance > best_op["performance"]
                  )
                  or (
                      self._optimize_mode is OptimizeMode.Minimize
                      and performance < best_op["performance"]
                  )
              ):
                  _logger.debug(
                      "updating best layer to %s...", wrapper.name
                  )
                  # find weight mask of this layer
                  for w in pruner.get_modules_wrapper():
                      if w.name == wrapper.name:
                          masks = {
                              "weight_mask": w.weight_mask,
                              "bias_mask": w.bias_mask,
                          }
                          break
                  best_op = {
                      "op_name": wrapper.name,
                      "sparsity": target_op_sparsity,
                      "performance": performance,
                      "masks": masks,
                  }
  
                  # save model weights
                  pruner.export_model(self._tmp_model_path)
  
          if not best_op:
              # decrease pruning step
              self._sparsity_per_iteration *= 0.5
              _logger.info(
                  "No more layers to prune, decrease pruning step to %s",
                  self._sparsity_per_iteration,
              )
              continue
  
          # Pick the best layer to prune, update iterative information
          # update config_list
          self._config_list_generated = self._update_config_list(
              self._config_list_generated,
              best_op["op_name"],
              best_op["sparsity"],
          )
  
          # update weights parameters
          self._model_to_prune.load_state_dict(
              torch.load(self._tmp_model_path)
          )
  
          # update mask of the chosen op
          for wrapper in self.get_modules_wrapper():
              if wrapper.name == best_op["op_name"]:
                  for k in best_op["masks"]:
                      setattr(wrapper, k, best_op["masks"][k])
                  break
  
          current_sparsity = target_sparsity
          _logger.info(
              "Pruning iteration %d finished, current sparsity: %s",
              pruning_iteration,
              current_sparsity,
          )
          _logger.info(
              "Layer %s seleted with sparsity %s, performance after pruning & short term fine-tuning : %s",
              best_op["op_name"],
              best_op["sparsity"],
              best_op["performance"],
          )
          pruning_iteration += 1
  
          self._final_performance = best_op["performance"]
  
      # load weights parameters
      self.load_model_state_dict(torch.load(self._tmp_model_path))
      os.remove(self._tmp_model_path)
  
      _logger.info("----------Compression finished--------------")
      _logger.info(
          "config_list generated: %s", self._config_list_generated
      )
      _logger.info(
          "Performance after pruning: %s", self._final_performance
      )
      _logger.info("Masked sparsity: %.6f", current_sparsity)
  
      # save best config found and best performance
      with open(
          os.path.join(self._experiment_data_dir, "search_result.json"),
          "w",
      ) as jsonfile:
          json.dump(
              {
                  "performance": self._final_performance,
                  "config_list": json.dumps(
                      self._config_list_generated
                  ),
              },
              jsonfile,
          )
  
      _logger.info(
          "search history and result saved to foler : %s",
          self._experiment_data_dir,
      )
  
      return self.bound_model
  
imports: import logging
  import os
  import copy
  import json
  import torch
  from schema import And, Optional
  from ..torch_utils import OptimizeMode
  from ..compressor import Pruner
  from ..utils.config_validation import CompressorSchema
  from ..utils.num_param_counter import get_total_num_weights
  from .constants_pruner import PRUNER_DICT
label_desc: Use image node_id 8 for calling the NetAdaptPruner obj's underlying member method code with example usage: obj.compress() and returns: self