apache_beam.ml.inference.onnx_inference module

class apache_beam.ml.inference.onnx_inference.OnnxModelHandlerNumpy(model_uri: str, session_options=None, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'], provider_options=None, *, inference_fn: Callable[[Sequence[numpy.ndarray], onnxruntime.capi.onnxruntime_inference_collection.InferenceSession, Optional[Dict[str, Any]]], Iterable[apache_beam.ml.inference.base.PredictionResult]] = <function default_numpy_inference_fn>)

Bases: apache_beam.ml.inference.base.ModelHandler

Implementation of the ModelHandler interface for onnx using numpy arrays as input. Note that inputs to ONNXModelHandler should be of the same sizes

Example Usage:

pcoll | RunInference(OnnxModelHandler(model_uri="my_uri"))

Parameters:

• model_uri – The URI to where the model is saved.
• inference_fn – The inference function to use on RunInference calls. default=default_numpy_inference_fn

load_model() → onnxruntime.capi.onnxruntime_inference_collection.InferenceSession

Loads and initializes an onnx inference session for processing.

run_inference(batch: Sequence[numpy.ndarray], inference_session: onnxruntime.capi.onnxruntime_inference_collection.InferenceSession, inference_args: Optional[Dict[str, Any]] = None) → Iterable[apache_beam.ml.inference.base.PredictionResult]

Runs inferences on a batch of numpy arrays.

Parameters:

• batch – A sequence of examples as numpy arrays. They should be single examples.
• inference_session – An onnx inference session. Must be runnable with input x where x is sequence of numpy array
• inference_args – Any additional arguments for an inference.

Returns:

An Iterable of type PredictionResult.

get_num_bytes(batch: Sequence[numpy.ndarray]) → int

Returns:The number of bytes of data for a batch.

get_metrics_namespace() → str

Returns:A namespace for metrics collected by the RunInference transform.