tfma.extractors.PredictExtractor

View source on GitHub

Creates an Extractor for TFMAPredict.

The extractor's PTransform loads and runs the eval_saved_model against every example yielding a copy of the Extracts input with an additional extract of type FeaturesPredictionsLabels keyed by tfma.FEATURES_PREDICTIONS_LABELS_KEY unless eval_config is not None in which case the features, predictions, and labels will be stored separately under tfma.FEATURES_KEY, tfma.PREDICTIONS_KEY, and tfma.LABELS_KEY respectively.

eval_shared_model Shared model (single-model evaluation) or list of shared models (multi-model evaluation).
desired_batch_size Optional batch size for batching in Aggregate.
materialize True to call the FeatureExtractor to add MaterializedColumn entries for the features, predictions, and labels.
eval_config Eval config.

Extractor for extracting features, predictions, labels, and other tensors during predict.