sparseml.onnx.sparsification package


sparseml.onnx.sparsification.analyzer module

Analyzer class implementations for ONNX

class sparseml.onnx.sparsification.analyzer.PruningLossSensitivityMagnitudeAnalyzer(model_info: sparseml.sparsification.model_info.ModelInfo)[source]

Bases: sparseml.sparsification.analyzer.PruningLossSensitivityMagnitudeAnalyzer

Class for performing weight mangitude pruning sensitivity analysis on ONNX models

pruning_loss_analysis_sparsity_levels is an optional run argument to set the sparsities that this analysis will run at. if not set, the value defaults to sparsml.optim.default_pruning_sparsities_loss(extended=True)

get_named_prunable_params(model: Any)Dict[str, numpy.ndarray][source]

loads the prunable parameters in a standardized way so that weight magnitude analysis may be run on each


model – model to load the prunable parameters from


dictionary of prunable parameter name as listed in the ModelInfo to a numpy array of the values of the parameter

static validate_model(prunable_param_names: Set[str], model: onnx.onnx_ml_pb2.ModelProto)bool[source]

Validates that all prunable parameter names in the ModelInfo layer_info exist in the given model and that the given model is of the correct framework

  • prunable_param_names – set of prunable parameter names found in the model info

  • model – model to validate


True if this is a valid model for weight mangitude pruning analysis. False otherwise

class sparseml.onnx.sparsification.analyzer.PruningPerformanceSensitivityAnalyzer(model_info: sparseml.sparsification.model_info.ModelInfo, batch_size: int = 1, num_cores: Optional[int] = None, iterations_per_check: int = 10, warmup_iterations_per_check: int = 5)[source]

Bases: sparseml.sparsification.analyzer.Analyzer

Class for running pruning performance sensitivity analysis on a model against the DeepSparse engine. deepsparse must be installed to be available.

pruning_perf_analysis_sparsity_levels is an optional run argument to set the sparisities that this analysis will run at. if not set, the value defaults to sparsml.optim.default_pruning_sparsities_perf()

  • model_info – ModelInfo object of the model to be analyzed. after running this analysis, the analysis_results of this ModelInfo object will be updated

  • batch_size – batch size to run analysis at. Default is 1

  • num_cores – number of CPU cores to run analysis with. Default is all available on the system

  • iterations_per_check – number of benchmarking iterations to run for each sparsity level. Default is 10

  • warmup_iterations_per_check – number of warmup iterations to run at each saprsity level. Default is 5

classmethod available(model_info: sparseml.sparsification.model_info.ModelInfo, **kwargs)bool[source]
  • model_info – ModelInfo object of the model to be analyzed

  • kwargs – keyword arguments that will be passed in to this analysis. model must be included for this analysis to be available


True if given the inputs, this analyzer can run its analysis. False otherwise


list of ONNX Analyzer implementations module

Functionality related to describing availability and information of sparsification algorithms to models within in the ONNX/ONNXRuntime framework.[source]

Load the available setup for sparsifying model within onnx.


The sparsification info for the onnx framework

Return type


sparseml.onnx.sparsification.model_info module

Classes for describing models and layers in ONNX models.

class sparseml.onnx.sparsification.model_info.ModelInfo(model: Any, metadata: Optional[Dict[str, Any]] = None)[source]

Bases: sparseml.sparsification.model_info.ModelInfo

Class for extracting and serializing ONNX model metadata, layers info, and analysis results


model – ONNX ModelProto object, or path to ONNX model file

extract_layer_info(model: onnx.onnx_ml_pb2.ModelProto)OrderedDict[str, LayerInfo][source]

model – ONNX model to extract LayerInfo of


ordered dictionary of layer name to LayerInfo object for the prunable model layers

static validate_model(model: Union[onnx.onnx_ml_pb2.ModelProto, str])[source]

Module contents

Functionality related to applying, describing, and supporting sparsification algorithms to models within in the ONNX/ONNXRuntime framework.