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