hyrax.verbs.engine

Attributes

logger

Classes

Engine

This verb drives inference with an ONNX model in production.

Module Contents

logger

class Engine(config)

Bases: hyrax.verbs.verb_registry.Verb

This verb drives inference with an ONNX model in production.

__init__()

Overall initialization for all verbs that saves the config

cli_name = 'engine'

add_parser_kwargs

description = 'Run inference with an ONNX model.'

static setup_parser(parser)

Setup parser for engine verb

run_cli(args=None)

CLI stub for Engine verb

run(model_directory: str = None)

Run inference with an ONNX model.

This method performs the following steps: - Read in the user config - Prepare all the datasets requested - Implement a simple strategy for reading in batches of data samples - Process the samples with any custom collate functions as well as a default collate function - Pass the collated batch to the appropriate to_tensor function - Send that output to the ONNX-ified model - Persist the results of inference

Parameters:

model_directory (str, optional) – Directory containing the ONNX model. If not provided, uses the config file or finds the most recent ONNX export directory.

create_ort_inputs(prepared_batch)

Create the inputs array for the ONNX model using the expected inputs from the loaded ONNX model and the type and shape of the prepared batch.

run_onnx_batch(ort_inputs)

Run the batch using our onnx runtime session

Only split out because this is when data is mutated and we need to be able to trace it.

_setup_trace(prepare_inputs_fn)