tfma.extractors.PredictionsExtractor

Stay organized with collections Save and categorize content based on your preferences.

Creates an extractor for performing predictions over a batch.

The extractor's PTransform loads and runs the serving saved_model(s) against every Extracts yielding a copy of the incoming Extracts with an additional Extracts added for the predictions keyed by tfma.PREDICTIONS_KEY. The model inputs are searched for under tfma.FEATURES_KEY (keras only) or tfma.INPUT_KEY (if tfma.FEATURES_KEY is not set or the model is non-keras). If multiple models are used the predictions will be stored in a dict keyed by model name.

Note that the prediction_key in the ModelSpecs also serves as a key into the dict of the prediction's output.

eval_config Eval config.
eval_shared_model Shared model (single-model evaluation) or list of shared models (multi-model evaluation) or None (predictions obtained from features).
experimental_bulk_inference Controls which inference implementation will be used. If True, will use the experimental TFX-BSL Bulk Inference implementation.
batch_size For testing only. Allows users to set a static batch size in unit tests.

Extractor for extracting predictions.