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)[source]

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[source]

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]][source]

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[source]

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[source]

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.