sparseml.onnx.sparsification package

Submodules

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

Parameters

model – model to load the prunable parameters from

Returns

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

Parameters

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

  • model – model to validate

Returns

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()

Parameters

  • 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]
Parameters

  • 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

Returns

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

sparseml.onnx.sparsification.analyzer.get_analyzer_impls() → List[sparseml.sparsification.analyzer.Analyzer][source]
Returns

list of ONNX Analyzer implementations

sparseml.onnx.sparsification.info module

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

sparseml.onnx.sparsification.info.sparsification_info() → sparseml.sparsification.info.SparsificationInfo[source]

Load the available setup for sparsifying model within onnx.

Returns

The sparsification info for the onnx framework

Return type

SparsificationInfo

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

Parameters

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

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

model – ONNX model to extract LayerInfo of

Returns

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.