def testMakeSklearnPredictExtractor(self):
"""Tests that predictions are made from extracts for a single model."""
feature_extractor = tfma.extractors.FeaturesExtractor(
self._eval_config)
prediction_extractor = (
sklearn_predict_extractor._make_sklearn_predict_extractor(
self._eval_shared_model))
with beam.Pipeline() as pipeline:
predict_extracts = (
pipeline
| 'Create' >> beam.Create(
[e.SerializeToString() for e in self._examples])
| 'BatchExamples' >> self._tfx_io.BeamSource()
| 'InputsToExtracts' >> tfma.BatchedInputsToExtracts() # pylint: disable=no-value-for-parameter
| feature_extractor.stage_name >> feature_extractor.ptransform
| prediction_extractor.stage_name >>
prediction_extractor.ptransform)
def check_result(actual):
try:
for item in actual:
self.assertEqual(item['labels'].shape,
item['predictions'].shape)
except AssertionError as err:
raise util.BeamAssertException(err)
util.assert_that(predict_extracts, check_result)
def benchmarkMiniPipeline(self):
"""Benchmark a "mini" TFMA - predict, slice and compute metrics.
Runs a "mini" version of TFMA in a Beam pipeline. Records the wall time
taken for the whole pipeline.
"""
self._init_model()
pipeline = self._create_beam_pipeline()
tfx_io = test_util.InMemoryTFExampleRecord(
schema=benchmark_utils.read_schema(
self._dataset.tf_metadata_schema_path()),
raw_record_column_name=constants.ARROW_INPUT_COLUMN)
raw_data = (
pipeline
| "Examples" >> beam.Create(
self._dataset.read_raw_dataset(
deserialize=False, limit=MAX_NUM_EXAMPLES))
| "BatchExamples" >> tfx_io.BeamSource()
| "InputsToExtracts" >> tfma.BatchedInputsToExtracts())
_ = (
raw_data
| "FeaturesExtractor" >> features_extractor.FeaturesExtractor(
eval_config=self._eval_config).ptransform
| "LabelsExtractor" >> labels_extractor.LabelsExtractor(
eval_config=self._eval_config).ptransform
| "ExampleWeightsExtractor" >> example_weights_extractor
.ExampleWeightsExtractor(eval_config=self._eval_config).ptransform
| "PredictionsExtractor" >> predictions_extractor.PredictionsExtractor(
eval_config=self._eval_config,
eval_shared_model=self._eval_shared_model).ptransform
| "UnbatchExtractor" >> unbatch_extractor.UnbatchExtractor().ptransform
| "SliceKeyExtractor" >> tfma.extractors.SliceKeyExtractor().ptransform
| "ComputeMetricsPlotsAndValidations" >>
metrics_plots_and_validations_evaluator
.MetricsPlotsAndValidationsEvaluator(
eval_config=self._eval_config,
eval_shared_model=self._eval_shared_model).ptransform)
start = time.time()
result = pipeline.run()
result.wait_until_finish()
end = time.time()
delta = end - start
self.report_benchmark(
iters=1,
wall_time=delta,
extras={
"num_examples": self._dataset.num_examples(limit=MAX_NUM_EXAMPLES)
})
def testMakeSklearnPredictExtractorWithMultiModels(self):
"""Tests that predictions are made from extracts for multiple models."""
eval_config = tfma.EvalConfig(model_specs=[
tfma.ModelSpec(name='model1'),
tfma.ModelSpec(name='model2'),
])
eval_export_dir_1 = os.path.join(self._eval_export_dir, '1')
self._create_sklearn_model(eval_export_dir_1)
eval_shared_model_1 = sklearn_predict_extractor.custom_eval_shared_model(
eval_saved_model_path=eval_export_dir_1,
model_name='model1',
eval_config=eval_config)
eval_export_dir_2 = os.path.join(self._eval_export_dir, '2')
self._create_sklearn_model(eval_export_dir_2)
eval_shared_model_2 = sklearn_predict_extractor.custom_eval_shared_model(
eval_saved_model_path=eval_export_dir_2,
model_name='model2',
eval_config=eval_config)
feature_extractor = tfma.extractors.FeaturesExtractor(
self._eval_config)
prediction_extractor = (
sklearn_predict_extractor._make_sklearn_predict_extractor(
eval_shared_model={
'model1': eval_shared_model_1,
'model2': eval_shared_model_2,
}))
with beam.Pipeline() as pipeline:
predict_extracts = (
pipeline
| 'Create' >> beam.Create(
[e.SerializeToString() for e in self._examples])
| 'BatchExamples' >> self._tfx_io.BeamSource()
| 'InputsToExtracts' >> tfma.BatchedInputsToExtracts() # pylint: disable=no-value-for-parameter
| feature_extractor.stage_name >> feature_extractor.ptransform
| prediction_extractor.stage_name >>
prediction_extractor.ptransform)
def check_result(actual):
try:
for item in actual:
self.assertEqual(item['labels'].shape,
item['predictions'].shape)
self.assertIn('model1', item['predictions'][0])
self.assertIn('model2', item['predictions'][0])
except AssertionError as err:
raise util.BeamAssertException(err)
util.assert_that(predict_extracts, check_result)
def benchmarkMiniPipelineBatched(self):
"""Benchmark a batched "mini" TFMA - predict, slice and compute metrics.
Runs a "mini" version of TFMA in a Beam pipeline. Records the wall time
taken for the whole pipeline.
"""
self._init_model()
pipeline = beam.Pipeline(runner=fn_api_runner.FnApiRunner())
tfx_io = test_util.InMemoryTFExampleRecord(
schema=benchmark_utils.read_schema(
self._dataset.tf_metadata_schema_path()),
raw_record_column_name=tfma.BATCHED_INPUT_KEY)
raw_data = (
pipeline
| "Examples" >> beam.Create(
self._dataset.read_raw_dataset(
deserialize=False, limit=MAX_NUM_EXAMPLES))
| "BatchExamples" >> tfx_io.BeamSource()
| "InputsToExtracts" >> tfma.BatchedInputsToExtracts())
_ = (
raw_data
| "BatchedInputExtractor" >> batched_input_extractor
.BatchedInputExtractor(eval_config=self._eval_config).ptransform
| "V2BatchedPredictExtractor" >>
batched_predict_extractor_v2.BatchedPredictExtractor(
eval_config=self._eval_config,
eval_shared_model=self._eval_shared_model).ptransform
| "UnbatchExtractor" >> unbatch_extractor.UnbatchExtractor().ptransform
| "SliceKeyExtractor" >> tfma.extractors.SliceKeyExtractor().ptransform
| "V2ComputeMetricsAndPlots" >>
metrics_and_plots_evaluator_v2.MetricsAndPlotsEvaluator(
eval_config=self._eval_config,
eval_shared_model=self._eval_shared_model).ptransform)
start = time.time()
result = pipeline.run()
result.wait_until_finish()
end = time.time()
delta = end - start
self.report_benchmark(
iters=1,
wall_time=delta,
extras={
"num_examples": self._dataset.num_examples(limit=MAX_NUM_EXAMPLES)
})
def _runMiniPipeline(self, multi_model):
"""Benchmark a "mini" TFMA - predict, slice and compute metrics.
Runs a "mini" version of TFMA in a Beam pipeline. Records the wall time
taken for the whole pipeline.
Args:
multi_model: True if multiple models should be used in the benchmark.
"""
self._init_model(multi_model, validation=False)
pipeline = self._create_beam_pipeline()
tfx_io = test_util.InMemoryTFExampleRecord(
schema=benchmark_utils.read_schema(
self._dataset.tf_metadata_schema_path()),
raw_record_column_name=constants.ARROW_INPUT_COLUMN)
raw_data = (pipeline
| "Examples" >> beam.Create(
self._dataset.read_raw_dataset(
deserialize=False, limit=self._max_num_examples()))
| "BatchExamples" >> tfx_io.BeamSource()
| "InputsToExtracts" >> tfma.BatchedInputsToExtracts())
def rescale_labels(extracts):
# Transform labels to [0, 1] so we can test metrics that require labels in
# that range.
result = copy.copy(extracts)
result[constants.LABELS_KEY] = self._transform_labels(
extracts[constants.LABELS_KEY])
return result
_ = (raw_data
| "FeaturesExtractor" >> features_extractor.FeaturesExtractor(
eval_config=self._eval_config).ptransform
| "LabelsExtractor" >> labels_extractor.LabelsExtractor(
eval_config=self._eval_config).ptransform
| "RescaleLabels" >> beam.Map(rescale_labels)
| "ExampleWeightsExtractor" >> example_weights_extractor.
ExampleWeightsExtractor(eval_config=self._eval_config).ptransform
| "PredictionsExtractor" >>
predictions_extractor.PredictionsExtractor(
eval_config=self._eval_config,
eval_shared_model=self._eval_shared_models).ptransform
| "UnbatchExtractor" >>
unbatch_extractor.UnbatchExtractor().ptransform
| "SliceKeyExtractor" >>
tfma.extractors.SliceKeyExtractor().ptransform
| "ComputeMetricsPlotsAndValidations" >>
metrics_plots_and_validations_evaluator.
MetricsPlotsAndValidationsEvaluator(
eval_config=self._eval_config,
eval_shared_model=self._eval_shared_models).ptransform)
start = time.time()
for _ in range(_ITERS):
result = pipeline.run()
result.wait_until_finish()
end = time.time()
delta = end - start
self.report_benchmark(
iters=_ITERS,
wall_time=delta,
extras={
"num_examples":
self._dataset.num_examples(limit=self._max_num_examples())
})