class_docstr: A Pytorch implementation of Simulated Annealing 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.
  evaluator : function
      Function to evaluate the pruned 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.
  start_temperature : float
      Start temperature of the simulated annealing process.
  stop_temperature : float
      Stop temperature of the simulated annealing process.
  cool_down_rate : float
      Cool down rate of the temperature.
  perturbation_magnitude : float
      Initial perturbation magnitude to the sparsities. The magnitude decreases with current temperature.
  experiment_data_dir : string
      PATH to save experiment data,
      including the config_list generated for the base pruning algorithm, the performance of the pruned model and the pruning history.
class_docstr_tok: ["A","Pytorch","implementation","of","Simulated","Annealing","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",".","evaluator",":","function","Function","to","evaluate","the","pruned","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",".","start_temperature",":","float","Start","temperature","of","the","simulated","annealing","process",".","stop_temperature",":","float","Stop","temperature","of","the","simulated","annealing","process",".","cool_down_rate",":","float","Cool","down","rate","of","the","temperature",".","perturbation_magnitude",":","float","Initial","perturbation","magnitude","to","the","sparsities",".","The","magnitude","decreases","with","current","temperature",".","experiment_data_dir",":","string","PATH","to","save","experiment","data",",","including","the","config_list","generated","for","the","base","pruning","algorithm",",","the","performance","of","the","pruned","model","and","the","pruning","history","."]
docstr: Compress the model with Simulated Annealing.
  
  Returns
  -------
  torch.nn.Module
      model with specified modules compressed.
docstr_tok: ["Compress","the","model","with","Simulated","Annealing",".","Returns","--","--","--","-","torch.nn.Module","model","with","specified","modules","compressed","."]
code: def compress(self, return_config_list=False):
      """
      Compress the model with Simulated Annealing.
  
      Returns
      -------
      torch.nn.Module
          model with specified modules compressed.
      """
      _logger.info("Starting Simulated Annealing Compression...")
  
      # initiaze a randomized action
      pruning_iteration = 0
      self._init_sparsities()
  
      # stop condition
      self._current_temperature = self._start_temperature
      while self._current_temperature > self._stop_temperature:
          _logger.info("Pruning iteration: %d", pruning_iteration)
          _logger.info(
              "Current temperature: %d, Stop temperature: %d",
              self._current_temperature,
              self._stop_temperature,
          )
          while True:
              # generate perturbation
              sparsities_perturbated = self._generate_perturbations()
              config_list = self._sparsities_2_config_list(
                  sparsities_perturbated
              )
              _logger.info(
                  "config_list for Pruner generated: %s", config_list
              )
  
              # fast evaluation
              pruner = PRUNER_DICT[self._base_algo](
                  copy.deepcopy(self._model_to_prune), config_list
              )
              model_masked = pruner.compress()
              evaluation_result = self._evaluator(model_masked)
  
              self._search_history.append(
                  {
                      "sparsity": self._sparsity,
                      "performance": evaluation_result,
                      "config_list": config_list,
                  }
              )
  
              if self._optimize_mode is OptimizeMode.Minimize:
                  evaluation_result *= -1
  
              # if better evaluation result, then accept the perturbation
              if evaluation_result > self._current_performance:
                  self._current_performance = evaluation_result
                  self._sparsities = sparsities_perturbated
  
                  # save best performance and best params
                  if evaluation_result > self._best_performance:
                      _logger.info("updating best model...")
                      self._best_performance = evaluation_result
                      self._best_config_list = config_list
  
                      # save the overall best masked model
                      self.bound_model = model_masked
                      # the ops with sparsity 0 are not included in this modules_wrapper
                      modules_wrapper_final = (
                          pruner.get_modules_wrapper()
                      )
                  break
              # if not, accept with probability e^(-deltaE/current_temperature)
              else:
                  delta_E = np.abs(
                      evaluation_result - self._current_performance
                  )
                  probability = math.exp(
                      -1 * delta_E / self._current_temperature
                  )
                  if np.random.uniform(0, 1) < probability:
                      self._current_performance = evaluation_result
                      self._sparsities = sparsities_perturbated
                      break
  
          # cool down
          self._current_temperature *= self._cool_down_rate
          pruning_iteration += 1
  
      _logger.info("----------Compression finished--------------")
      _logger.info("Best performance: %s", self._best_performance)
      _logger.info("config_list found : %s", self._best_config_list)
  
      # save search history
      with open(
          os.path.join(self._experiment_data_dir, "search_history.csv"),
          "w",
      ) as csvfile:
          writer = csv.DictWriter(
              csvfile,
              fieldnames=["sparsity", "performance", "config_list"],
          )
          writer.writeheader()
          for item in self._search_history:
              writer.writerow(
                  {
                      "sparsity": item["sparsity"],
                      "performance": item["performance"],
                      "config_list": json.dumps(item["config_list"]),
                  }
              )
  
      # save best config found and best performance
      if self._optimize_mode is OptimizeMode.Minimize:
          self._best_performance *= -1
      with open(
          os.path.join(self._experiment_data_dir, "search_result.json"),
          "w+",
      ) as jsonfile:
          json.dump(
              {
                  "performance": self._best_performance,
                  "config_list": json.dumps(self._best_config_list),
              },
              jsonfile,
          )
  
      _logger.info(
          "search history and result saved to foler : %s",
          self._experiment_data_dir,
      )
  
      if return_config_list:
          return self._best_config_list
  
      # This should be done only at the final stage,
      # because the modules_wrapper with all the ops are used during the annealing process
      self.modules_wrapper = modules_wrapper_final
  
      return self.bound_model
  
imports: import logging
  import os
  import math
  import copy
  import csv
  import json
  import numpy
  from schema import And, Optional
  from ..torch_utils import OptimizeMode
  from ..compressor import Pruner
  from ..utils.config_validation import CompressorSchema
  from .constants_pruner import PRUNER_DICT
label_desc: Use image node_id 8 for calling the SimulatedAnnealingPruner obj's underlying member method code with example usage: obj.compress(return_config_list) and returns: self, self