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[[tensorflow.python.module.module.Module, Sequence[Union[numpy.ndarray, tensorflow.python.framework.ops.Tensor]], 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() → tensorflow.python.module.module.Module[source]

Loads and initializes a Tensorflow model for processing.

update_model_path(model_path: Optional[str] = None)[source]
run_inference(batch: Sequence[numpy.ndarray], model: tensorflow.python.module.module.Module, 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.
get_metrics_namespace() → str[source]
Returns:A namespace for metrics collected by the RunInference transform.
validate_inference_args(inference_args: Optional[Dict[str, Any]])[source]
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[[tensorflow.python.module.module.Module, Sequence[Union[numpy.ndarray, tensorflow.python.framework.ops.Tensor]], 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() → tensorflow.python.module.module.Module[source]

Loads and initializes a tensorflow model for processing.

update_model_path(model_path: Optional[str] = None)[source]
run_inference(batch: Sequence[tensorflow.python.framework.ops.Tensor], model: tensorflow.python.module.module.Module, 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[tensorflow.python.framework.ops.Tensor]) → int[source]
Returns:The number of bytes of data for a batch of Tensors.
get_metrics_namespace() → str[source]
Returns:A namespace for metrics collected by the RunInference transform.
validate_inference_args(inference_args: Optional[Dict[str, Any]])[source]