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()) })