class_docstr: A Pytorch implementation of AutoCompress pruning 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.
  trainer : function
      Function used for the first subproblem of ADMM Pruner.
      Users should write this function as a normal function to train the Pytorch model
      and include `model, optimizer, criterion, epoch, callback` as function arguments.
      Here `callback` acts as an L2 regulizer as presented in the formula (7) of the original paper.
      The logic of `callback` is implemented inside the Pruner,
      users are just required to insert `callback()` between `loss.backward()` and `optimizer.step()`.
      Example::
  
          def trainer(model, criterion, optimizer, epoch, callback):
              device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
              train_loader = ...
              model.train()
              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()
                  # callback should be inserted between loss.backward() and optimizer.step()
                  if callback:
                      callback()
                  optimizer.step()
  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
  dummy_input : pytorch tensor
      The dummy input for ```jit.trace```, users should put it on right device before pass in.
  num_iterations : int
      Number of overall iterations.
  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.
  admm_num_iterations : int
      Number of iterations of ADMM Pruner.
  admm_training_epochs : int
      Training epochs of the first optimization subproblem of ADMMPruner.
  row : float
      Penalty parameters for ADMM training.
  experiment_data_dir : string
      PATH to store temporary experiment data.
class_docstr_tok: ["A","Pytorch","implementation","of","AutoCompress","pruning","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",".","trainer",":","function","Function","used","for","the","first","subproblem","of","ADMM","Pruner",".","Users","should","write","this","function","as","a","normal","function","to","train","the","Pytorch","model","and","include","`","model",",","optimizer",",","criterion",",","epoch",",","callback","`","as","function","arguments",".","Here","`","callback","`","acts","as","an","L2","regulizer","as","presented","in","the","formula","(","7",")","of","the","original","paper",".","The","logic","of","`","callback","`","is","implemented","inside","the","Pruner",",","users","are","just","required","to","insert","`","callback","(",")","`","between","`","loss.backward","(",")","`","and","`","optimizer.step","(",")","`",".","Example",":",":","def","trainer","(","model",",","criterion",",","optimizer",",","epoch",",","callback",")",":","device","=","torch.device","(","``","cuda","''","if","torch.cuda.is_available","(",")","else","``","cpu","''",")","train_loader","=","...","model.train","(",")","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","(",")","#","callback","should","be","inserted","between","loss.backward","(",")","and","optimizer.step","(",")","if","callback",":","callback","(",")","optimizer.step","(",")","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","dummy_input",":","pytorch","tensor","The","dummy","input","for","``","`","jit.trace","``","`",",","users","should","put","it","on","right","device","before","pass","in",".","num_iterations",":","int","Number","of","overall","iterations",".","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",".","admm_num_iterations",":","int","Number","of","iterations","of","ADMM","Pruner",".","admm_training_epochs",":","int","Training","epochs","of","the","first","optimization","subproblem","of","ADMMPruner",".","row",":","float","Penalty","parameters","for","ADMM","training",".","experiment_data_dir",":","string","PATH","to","store","temporary","experiment","data","."]
code: def export_model(
      self,
      model_path,
      mask_path=None,
      onnx_path=None,
      input_shape=None,
      device=None,
  ):
      _logger.info(
          "AutoCompressPruner export directly the pruned model without mask"
      )
  
      torch.save(self._model_to_prune.state_dict(), model_path)
      _logger.info("Model state_dict saved to %s", model_path)
  
      if onnx_path is not None:
          assert (
              input_shape is not None
          ), "input_shape must be specified to export onnx model"
          # input info needed
          if device is None:
              device = torch.device("cpu")
          input_data = torch.Tensor(*input_shape)
          torch.onnx.export(
              self._model_to_prune, input_data.to(device), onnx_path
          )
          _logger.info(
              "Model in onnx with input shape %s saved to %s",
              input_data.shape,
              onnx_path,
          )
  
file_path: auptimizer/src/aup/compression/torch/pruning/auto_compress_pruner.py
imports: import logging
  import os
  import copy
  import torch
  from schema import And, Optional
  from ..torch_utils import OptimizeMode
  from ..None import ModelSpeedup
  from ..compressor import Pruner
  from ..utils.config_validation import CompressorSchema
  from .simulated_annealing_pruner import SimulatedAnnealingPruner
  from .admm_pruner import ADMMPruner
label_desc: Use image node_id 5 for calling the AutoCompressPruner obj's underlying member method code with example usage: obj.export_model(model_path, mask_path, onnx_path, input_shape, device) without return types