sparseml.onnx.benchmark package

Submodules

sparseml.onnx.benchmark.info module

class sparseml.onnx.benchmark.info.ORTBenchmarkRunner(model: Any, batch_size: int = 1, iterations: int = 0, warmup_iterations: int = 0, framework_args: Dict[str, Any] = {}, provider: str = 'cpu', device: str = 'cpu', ort_provider: Optional[str] = None, **kwargs)

Bases: sparseml.benchmark.info.BenchmarkRunner

Benchmark runner for ONNXruntime.

Parameters

• model – model to benchmark

• batch_size – batch size to use for benchmarking

• iterations – number of iterations to run

• warmup_iterations – number of warmup iterations to run

• framework_args – additional arguments to pass to the framework

• provider – inference provider name to use from available FrameworkInfo

• device – device to use for benchmarking

• ort_provider – provider to use for ONNXruntime

property batch_size

the batch size

Type

return

property device

the device

Type

return

property framework

the framework

Type

return

property framework_args

the framework args

Type

return

property framework_info

the framework info

Type

return

property inference_provider

the inference provider

Type

return

property iterations

the number of iterations

Type

return

property model

the model as an ONNX ModelProto

Type

return

property num_cores

the number of cores

Type

return

property package_versions

the package versions

Type

return

run_batch(batch: Union[Dict[str, Any], Tuple[Dict[str, Any], Any]], *args, **kwargs) → sparseml.benchmark.serialization.BatchBenchmarkResult

Runs a benchmark on a given batch.

Parameters

• batch – the batch to benchmark

• args – additional arguments to pass to the framework

• kwargs – additional arguments to pass to the framework

property warmup_iterations

the warmup iterations

Type

return

sparseml.onnx.benchmark.info.load_data(data: Any, model: Optional[Any] = None, batch_size: int = 1, total_iterations: int = 0, **kwargs) → Iterable[Tuple[Dict[str, Any], Any]]

Creates a iteratable data loader for the given data.

Acceptable types for data are: - a folder path containing numpy files - a list of file paths - a SparseML DataLoader - a SparseZoo DataLoader - an iterable - None type, in which case model must be passed

Parameters

• data – data to use for benchmarking

• model – model to use for generating data

• batch_size – batch size

• total_iterations – total number of iterations

• kwargs – additional arguments to pass to the DataLoader

Returns

an iterable of data and labels

sparseml.onnx.benchmark.info.load_model(model: Any, **kwargs) → onnx.onnx_ml_pb2.ModelProto

Loads the model and saves it to a temporary file if necessary

Parameters

• model – the model

• kwargs – additional arguments to pass if loading from a stub

Returns

the model loaded as a ModelProto

sparseml.onnx.benchmark.info.run_benchmark(model: Any, data: Any = None, batch_size: int = 1, iterations: int = 0, warmup_iterations: int = 0, provider: Optional[str] = 'cpu', device: Optional[str] = 'cpu', framework_args: Dict[str, Any] = {}, show_progress: bool = True, **kwargs) → sparseml.benchmark.serialization.BenchmarkInfo

Run a benchmark for the given model.

Parameters

• model – model to benchmark

• data – data to benchmark

• batch_size – batch size

• iterations – number of iterations

• warmup_iterations – number of warmup iterations

• framework – the specific framework run the benchmark in

• provider – the specific inference provider to use

• device – the specific device to use

• save_path – path to save the benchmark results

• framework_args – additional framework specific arguments to pass to the runner

• show_progress – True to show a tqdm bar when running, False otherwise

• kwargs – Additional arguments to pass to the framework.

Returns

BenchmarkInfo

Module contents

Functionality related to running benchmarks in the ONNX/ONNXruntime frameworks.