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], <sphinx.ext.autodoc.importer._MockObject object at 0x7ff1e191f160>, Optional[Dict[str, Any]]], Iterable[apache_beam.ml.inference.base.PredictionResult]] = <function default_numpy_inference_fn>, **kwargs)[source]¶ Bases:
apache_beam.ml.inference.base.ModelHandlerImplementation 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
- kwargs – ‘env_vars’ can be used to set environment variables before loading the model.
-
load_model() → <sphinx.ext.autodoc.importer._MockObject object at 0x7ff1df41aa90>[source]¶ Loads and initializes an onnx inference session for processing.
-
run_inference(batch: Sequence[numpy.ndarray], inference_session: <sphinx.ext.autodoc.importer._MockObject object at 0x7ff1df41a7c0>, 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.