apache_beam.ml.inference.tensorflow_inference module¶
-
class
apache_beam.ml.inference.tensorflow_inference.
TFModelHandlerNumpy
(model_uri: str, model_type: apache_beam.ml.inference.tensorflow_inference.ModelType = <ModelType.SAVED_MODEL: 1>, create_model_fn: Optional[Callable] = None, *, inference_fn: Callable[[<sphinx.ext.autodoc.importer._MockObject object at 0x7f2e2e484e20>, Sequence[Union[numpy.ndarray, <sphinx.ext.autodoc.importer._MockObject object at 0x7f2e2e484e80>]], Dict[str, Any], Optional[str]], 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 Tensorflow.
Example Usage:
pcoll | RunInference(TFModelHandlerNumpy(model_uri="my_uri"))
See https://www.tensorflow.org/tutorials/keras/save_and_load for details.
Parameters: - model_uri (str) – path to the trained model.
- model_type – type of model to be loaded. Defaults to SAVED_MODEL.
- create_model_fn – a function that creates and returns a new tensorflow model to load the saved weights. It should be used with ModelType.SAVED_WEIGHTS.
- inference_fn – inference function to use during RunInference. Defaults to default_numpy_inference_fn.
Supported Versions: RunInference APIs in Apache Beam have been tested with Tensorflow 2.9, 2.10, 2.11.
-
load_model
() → <sphinx.ext.autodoc.importer._MockObject object at 0x7f2e2e4841c0>[source]¶ Loads and initializes a Tensorflow model for processing.
-
run_inference
(batch: Sequence[numpy.ndarray], model: <sphinx.ext.autodoc.importer._MockObject object at 0x7f2e2e484070>, inference_args: Optional[Dict[str, Any]] = None) → Iterable[apache_beam.ml.inference.base.PredictionResult][source]¶ Runs inferences on a batch of numpy array and returns an Iterable of numpy array Predictions.
This method stacks the n-dimensional numpy array in a vectorized format to optimize the inference call.
Parameters: - batch – A sequence of numpy nd-array. These should be batchable, as this method will call numpy.stack() and pass in batched numpy nd-array with dimensions (batch_size, n_features, etc.) into the model’s predict() function.
- model – A Tensorflow model.
- inference_args – any additional arguments for an inference.
Returns: An Iterable of type PredictionResult.
-
get_num_bytes
(batch: Sequence[numpy.ndarray]) → int[source]¶ Returns: The number of bytes of data for a batch of numpy arrays.
-
class
apache_beam.ml.inference.tensorflow_inference.
TFModelHandlerTensor
(model_uri: str, model_type: apache_beam.ml.inference.tensorflow_inference.ModelType = <ModelType.SAVED_MODEL: 1>, create_model_fn: Optional[Callable] = None, *, inference_fn: Callable[[<sphinx.ext.autodoc.importer._MockObject object at 0x7f2e2e484e20>, Sequence[Union[numpy.ndarray, <sphinx.ext.autodoc.importer._MockObject object at 0x7f2e2e484e80>]], Dict[str, Any], Optional[str]], Iterable[apache_beam.ml.inference.base.PredictionResult]] = <function default_tensor_inference_fn>)[source]¶ Bases:
apache_beam.ml.inference.base.ModelHandler
Implementation of the ModelHandler interface for Tensorflow.
Example Usage:
pcoll | RunInference(TFModelHandlerTensor(model_uri="my_uri"))
See https://www.tensorflow.org/tutorials/keras/save_and_load for details.
Parameters: - model_uri (str) – path to the trained model.
- model_type – type of model to be loaded. Defaults to SAVED_MODEL.
- create_model_fn – a function that creates and returns a new tensorflow model to load the saved weights. It should be used with ModelType.SAVED_WEIGHTS.
- inference_fn – inference function to use during RunInference. Defaults to default_numpy_inference_fn.
Supported Versions: RunInference APIs in Apache Beam have been tested with Tensorflow 2.11.
-
load_model
() → <sphinx.ext.autodoc.importer._MockObject object at 0x7f2e2e52e7f0>[source]¶ Loads and initializes a tensorflow model for processing.
-
run_inference
(batch: Sequence[<sphinx.ext.autodoc.importer._MockObject object at 0x7f2e2e492100>], model: <sphinx.ext.autodoc.importer._MockObject object at 0x7f2e2e492160>, inference_args: Optional[Dict[str, Any]] = None) → Iterable[apache_beam.ml.inference.base.PredictionResult][source]¶ Runs inferences on a batch of tf.Tensor and returns an Iterable of Tensor Predictions.
This method stacks the list of Tensors in a vectorized format to optimize the inference call.
Parameters: - batch – A sequence of Tensors. These Tensors should be batchable, as this method will call tf.stack() and pass in batched Tensors with dimensions (batch_size, n_features, etc.) into the model’s predict() function.
- model – A Tensorflow model.
- inference_args – Non-batchable arguments required as inputs to the model’s forward() function. Unlike Tensors in batch, these parameters will not be dynamically batched
Returns: An Iterable of type PredictionResult.
-
get_num_bytes
(batch: Sequence[<sphinx.ext.autodoc.importer._MockObject object at 0x7f2e2e492250>]) → int[source]¶ Returns: The number of bytes of data for a batch of Tensors.