def _get_common_variables(dataset, force_tf_compat_v1): """Returns metadata schema, preprocessing fn, input dataset metadata.""" tf_metadata_schema = benchmark_utils.read_schema( dataset.tf_metadata_schema_path()) preprocessing_fn = dataset.tft_preprocessing_fn() feature_spec = schema_utils.schema_as_feature_spec( tf_metadata_schema).feature_spec type_spec = impl_helper.get_type_specs_from_feature_specs(feature_spec) transform_input_columns = ( tft.get_transform_input_columns( preprocessing_fn, type_spec, force_tf_compat_v1=force_tf_compat_v1)) transform_input_dataset_metadata = dataset_metadata.DatasetMetadata( schema_utils.schema_from_feature_spec({ feature: feature_spec[feature] for feature in transform_input_columns })) tfxio = tf_example_record.TFExampleBeamRecord( physical_format="tfexamples", schema=transform_input_dataset_metadata.schema, telemetry_descriptors=["TFTransformBenchmark"]) return CommonVariablesTuple( tf_metadata_schema=tf_metadata_schema, preprocessing_fn=preprocessing_fn, transform_input_dataset_metadata=transform_input_dataset_metadata, tfxio=tfxio)
def test_features_extractor_no_features(self): model_spec = config_pb2.ModelSpec() eval_config = config_pb2.EvalConfig(model_specs=[model_spec]) feature_extractor = features_extractor.FeaturesExtractor(eval_config) tfx_io = tf_example_record.TFExampleBeamRecord( raw_record_column_name=constants.ARROW_INPUT_COLUMN, physical_format='inmem', telemetry_descriptors=['testing']) with beam.Pipeline() as pipeline: result = ( pipeline | 'Create' >> beam.Create([b''] * 3) | 'DecodeToRecordBatch' >> tfx_io.BeamSource(batch_size=3) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | feature_extractor.stage_name >> feature_extractor.ptransform) def check_result(got): self.assertLen(got, 1) self.assertIn(constants.FEATURES_KEY, got[0]) self.assertEmpty(got[0][constants.FEATURES_KEY]) self.assertIn(constants.INPUT_KEY, got[0]) self.assertLen(got[0][constants.INPUT_KEY], 3) util.assert_that(result, check_result, label='CheckResult')
def testE2E(self): raw_record_column_name = "raw_record" tfxio = tf_example_record.TFExampleBeamRecord( physical_format="inmem", telemetry_descriptors=["some", "component"], schema=_SCHEMA, raw_record_column_name=raw_record_column_name, ) def _AssertFn(record_batches): self.assertLen(record_batches, 1) record_batch = record_batches[0] self.assertTrue(record_batch.schema.equals(tfxio.ArrowSchema())) tensor_adapter = tfxio.TensorAdapter() dict_of_tensors = tensor_adapter.ToBatchTensors(record_batch) self.assertLen(dict_of_tensors, 3) self.assertIn("int_feature", dict_of_tensors) self.assertIn("float_feature", dict_of_tensors) self.assertIn("string_feature", dict_of_tensors) with beam.Pipeline() as p: record_batch_pcoll = ( p | "CreateInMemRecords" >> beam.Create(_SERIALIZED_EXAMPLES) | "BeamSource" >> tfxio.BeamSource(batch_size=len(_SERIALIZED_EXAMPLES))) beam_testing_util.assert_that(record_batch_pcoll, _AssertFn)
def testModelSignaturesDoFn(self, save_as_keras, signature_names, default_signature_names, prefer_dict_outputs, use_schema, expected_num_outputs): export_path = self.createModelWithMultipleDenseInputs(save_as_keras) eval_shared_models = {} model_specs = [] for sigs in signature_names.values(): for model_name in sigs: if model_name not in eval_shared_models: eval_shared_models[ model_name] = self.createTestEvalSharedModel( eval_saved_model_path=export_path, model_name=model_name, tags=[tf.saved_model.SERVING]) model_specs.append(config_pb2.ModelSpec(name=model_name)) eval_config = config_pb2.EvalConfig(model_specs=model_specs) schema = self.createDenseInputsSchema() if use_schema else None tfx_io = tf_example_record.TFExampleBeamRecord( physical_format='text', schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN) examples = [ self._makeExample(input_1=1.0, input_2=2.0), self._makeExample(input_1=3.0, input_2=4.0), self._makeExample(input_1=5.0, input_2=6.0), ] with beam.Pipeline() as pipeline: # pylint: disable=no-value-for-parameter result = (pipeline | 'Create' >> beam.Create( [e.SerializeToString() for e in examples]) | 'BatchExamples' >> tfx_io.BeamSource(batch_size=3) | 'ToExtracts' >> beam.Map(_record_batch_to_extracts) | 'ModelSignatures' >> beam.ParDo( model_util.ModelSignaturesDoFn( eval_config=eval_config, eval_shared_models=eval_shared_models, signature_names=signature_names, default_signature_names=default_signature_names, prefer_dict_outputs=prefer_dict_outputs))) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) for key in signature_names: self.assertIn(key, got[0]) if prefer_dict_outputs: self.assertIsInstance(got[0][key], dict) self.assertEqual(tfma_util.batch_size(got[0][key]), expected_num_outputs) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
def testSqlSliceKeyExtractorWithEmptySqlConfig(self): eval_config = config_pb2.EvalConfig() feature_extractor = features_extractor.FeaturesExtractor( eval_config=eval_config) slice_key_extractor = sql_slice_key_extractor.SqlSliceKeyExtractor( eval_config) tfx_io = tf_example_record.TFExampleBeamRecord( physical_format='inmem', telemetry_descriptors=['test', 'component'], schema=_SCHEMA, raw_record_column_name=constants.ARROW_INPUT_COLUMN) examples = [ self._makeExample(fixed_int=1, fixed_float=1.0, fixed_string='fixed_string1'), self._makeExample(fixed_int=1, fixed_float=1.0, fixed_string='fixed_string2'), self._makeExample(fixed_int=2, fixed_float=0.0, fixed_string='fixed_string3') ] with beam.Pipeline() as pipeline: # pylint: disable=no-value-for-parameter result = ( pipeline | 'Create' >> beam.Create( [e.SerializeToString() for e in examples], reshuffle=False) | 'BatchExamples' >> tfx_io.BeamSource(batch_size=3) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | feature_extractor.stage_name >> feature_extractor.ptransform | slice_key_extractor.stage_name >> slice_key_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) np.testing.assert_equal( got[0][constants.SLICE_KEY_TYPES_KEY], types.VarLenTensorValue.from_dense_rows( [np.array([]), np.array([]), np.array([])])) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result)
def testSqlSliceKeyExtractor(self): eval_config = config_pb2.EvalConfig(slicing_specs=[ config_pb2.SlicingSpec(slice_keys_sql=""" SELECT STRUCT(fixed_string) FROM example.fixed_string, example.fixed_int WHERE fixed_int = 1 """) ]) slice_key_extractor = sql_slice_key_extractor.SqlSliceKeyExtractor( eval_config) tfx_io = tf_example_record.TFExampleBeamRecord( physical_format='inmem', telemetry_descriptors=['test', 'component'], schema=_SCHEMA, raw_record_column_name=constants.ARROW_INPUT_COLUMN) examples = [ self._makeExample( fixed_int=1, fixed_float=1.0, fixed_string='fixed_string1'), self._makeExample( fixed_int=1, fixed_float=1.0, fixed_string='fixed_string2'), self._makeExample( fixed_int=2, fixed_float=0.0, fixed_string='fixed_string3') ] with beam.Pipeline() as pipeline: # pylint: disable=no-value-for-parameter result = ( pipeline | 'Create' >> beam.Create([e.SerializeToString() for e in examples], reshuffle=False) | 'BatchExamples' >> tfx_io.BeamSource(batch_size=3) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | slice_key_extractor.stage_name >> slice_key_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) self.assertEqual(got[0][constants.SLICE_KEY_TYPES_KEY], [[(('fixed_string', 'fixed_string1'),)], [(('fixed_string', 'fixed_string2'),)], []]) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result)
def testModelSignaturesDoFnError(self): export_path = self.createModelWithInvalidOutputShape() signature_names = {constants.PREDICTIONS_KEY: {'': [None]}} eval_shared_models = { '': self.createTestEvalSharedModel(eval_saved_model_path=export_path, tags=[tf.saved_model.SERVING]) } model_specs = [config_pb2.ModelSpec()] eval_config = config_pb2.EvalConfig(model_specs=model_specs) schema = self.createDenseInputsSchema() tfx_io = tf_example_record.TFExampleBeamRecord( physical_format='text', schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN) tensor_adapter_config = tensor_adapter.TensorAdapterConfig( arrow_schema=tfx_io.ArrowSchema(), tensor_representations=tfx_io.TensorRepresentations()) examples = [ self._makeExample(input_1=1.0, input_2=2.0), self._makeExample(input_1=3.0, input_2=4.0), self._makeExample(input_1=5.0, input_2=6.0), ] with self.assertRaisesRegex( ValueError, 'First dimension does not correspond with batch size.'): with beam.Pipeline() as pipeline: # pylint: disable=no-value-for-parameter _ = (pipeline | 'Create' >> beam.Create( [e.SerializeToString() for e in examples]) | 'BatchExamples' >> tfx_io.BeamSource(batch_size=3) | 'ToExtracts' >> beam.Map(_record_batch_to_extracts) | 'ModelSignatures' >> beam.ParDo( model_util.ModelSignaturesDoFn( eval_config=eval_config, eval_shared_models=eval_shared_models, signature_names=signature_names, default_signature_names=None, prefer_dict_outputs=False, tensor_adapter_config=tensor_adapter_config)))
def compute_cis(scenario: str, methodology: str, num_trials: int, num_examples_per_trial: int, output_dir: str) -> None: """Computes a collection of CIs and the population values for a scenario.""" if scenario == _BINARY_CLASSIFICATION_SCENARIO: eval_config, example_gen_fn = get_binary_classification_scenario() elif scenario == _REGRESSION_SCENARIO: eval_config, example_gen_fn = get_regression_scenario() else: raise ValueError(f'Unexpected scenario {scenario}. ' f'Expected one of {_SCENARIOS}') eval_config.options.compute_confidence_intervals.value = True eval_config.options.confidence_intervals.method = ( config_pb2.ConfidenceIntervalOptions.ConfidenceIntervalMethod.Value( methodology)) pipeline_options = beam.options.pipeline_options.PipelineOptions( FLAGS.pipeline_options.split(',')) with beam.Pipeline(options=pipeline_options) as pipeline: tfx_io = tf_example_record.TFExampleBeamRecord( physical_format='generated', raw_record_column_name=constants.ARROW_INPUT_COLUMN) inputs_per_trial = [] for i in range(num_trials): inputs = (pipeline | f'CreateExamples[{i}]' >> beam.Create( example_gen_fn(num_examples_per_trial)) | f'Serialize[{i}]' >> beam.Map(lambda example: example.SerializeToString()) | f'BatchExamples[{i}]' >> tfx_io.BeamSource()) inputs_per_trial.append(inputs) trial_output_dir = os.path.join(output_dir, str(i)) _ = (inputs | f'Evaluate[{i}]' >> model_eval_lib.ExtractEvaluateAndWriteResults( eval_config=eval_config, output_path=trial_output_dir)) population_output_dir = os.path.join(output_dir, _POPULATION_OUTPUT_NAME) _ = (inputs_per_trial | 'FlattenInputs' >> beam.Flatten() | 'EvaluatePopulation' >> model_eval_lib.ExtractEvaluateAndWriteResults( eval_config=eval_config, output_path=population_output_dir))
def test_features_extractor(self): model_spec = config_pb2.ModelSpec() eval_config = config_pb2.EvalConfig(model_specs=[model_spec]) feature_extractor = features_extractor.FeaturesExtractor(eval_config) schema = text_format.Parse( """ feature { name: "example_weight" type: FLOAT } feature { name: "fixed_int" type: INT } feature { name: "fixed_float" type: FLOAT } feature { name: "fixed_string" type: BYTES } """, schema_pb2.Schema()) tfx_io = tf_example_record.TFExampleBeamRecord( schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN, physical_format='inmem', telemetry_descriptors=['testing']) example_kwargs = [ { 'fixed_int': 1, 'fixed_float': 1.0, 'fixed_string': 'fixed_string1' }, { 'fixed_int': 1, 'fixed_float': 1.0, 'fixed_string': 'fixed_string2' }, { 'fixed_int': 2, 'fixed_float': 0.0, 'fixed_string': 'fixed_string3' }, ] with beam.Pipeline() as pipeline: # pylint: disable=no-value-for-parameter result = ( pipeline | 'Create' >> beam.Create([ self._makeExample(**kwargs).SerializeToString() for kwargs in example_kwargs ], reshuffle=False) | 'DecodeToRecordBatch' >> tfx_io.BeamSource(batch_size=3) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | feature_extractor.stage_name >> feature_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) self.assertIn(constants.FEATURES_KEY, got[0]) self.assertLen(got[0][constants.FEATURES_KEY], 4) # 4 features self.assertIn('example_weight', got[0][constants.FEATURES_KEY]) # Arrays of type np.object won't compare with assertAllClose self.assertEqual( got[0][constants.FEATURES_KEY] ['example_weight'].tolist(), [None, None, None]) self.assertIn('fixed_int', got[0][constants.FEATURES_KEY]) self.assertAllClose( got[0][constants.FEATURES_KEY]['fixed_int'], np.array([1, 1, 2])) self.assertIn('fixed_float', got[0][constants.FEATURES_KEY]) self.assertAllClose( got[0][constants.FEATURES_KEY]['fixed_float'], np.array([1.0, 1.0, 0.0])) self.assertIn('fixed_string', got[0][constants.FEATURES_KEY]) # Arrays of type np.object won't compare with assertAllClose self.assertEqual( got[0][ constants.FEATURES_KEY]['fixed_string'].tolist(), [b'fixed_string1', b'fixed_string2', b'fixed_string3']) self.assertIn(constants.INPUT_KEY, got[0]) self.assertLen(got[0][constants.INPUT_KEY], 3) # 3 examples except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
def append_tfma_pipeline(pipeline: beam.Pipeline, me_eval_config: me_proto.EvaluationConfig, problem_type: constants.ProblemType, tfma_format: Optional[bool] = False, json_mode: Optional[bool] = False, schema: Optional[Any] = None): """Extend a beam pipeline to add TFMA evaluation given a configuration. Args: pipeline: A beam pipeline. me_eval_config: A ME Evaluation Configuration. problem_type: Defines what type of problem to expect. tfma_format: If true, use TFMA format, if false use Model Evaluation. json_mode: Output metrics in a plain text mode. schema: Optional tf.metadata schema. If you need to pass multi-tensor input to the model, you need to pass the schema. """ input_files = ( me_eval_config.data_spec.input_source_spec.jsonl_file_spec.file_names) output_path = me_eval_config.output_spec.gcs_sink.path data_spec = me_eval_config.data_spec weight_column_spec = ColumnSpec( me_eval_config.data_spec.example_weight_key_spec ) if me_eval_config.data_spec.HasField('example_weight_key_spec') else None eval_column_specs = EvaluationColumnSpecs( ground_truth_column_spec=ColumnSpec( me_eval_config.data_spec.label_key_spec), example_weight_column_spec=weight_column_spec, predicted_score_column_spec=ColumnSpec( data_spec.predicted_score_key_spec) if data_spec.HasField('predicted_score_key_spec') else None, predicted_label_column_spec=ColumnSpec( data_spec.predicted_label_key_spec) if data_spec.HasField('predicted_label_key_spec') else None, predicted_label_id_column_spec=ColumnSpec( data_spec.predicted_label_id_key_spec) if data_spec.HasField('predicted_label_id_key_spec') else None) class_name_list = list(me_eval_config.data_spec.labels) or None quantile_list = list(data_spec.quantiles) or None quantile_index = data_spec.quantile_index if data_spec.quantile_index >= 0 else None tfma_eval_config = tfma_adapter.METoTFMA(class_name_list).eval_config( me_eval_config) me_writers = [ tfma.writers.Writer( stage_name='WriteMetrics', # pylint:disable=no-value-for-parameter ptransform=_write_metrics(output_file=os.path.join( output_path, constants.Pipeline.METRICS_KEY), problem_type=problem_type, class_labels=class_name_list, tfma_format=tfma_format, json_mode=json_mode)), ] coder = tf_example_record.TFExampleBeamRecord( physical_format='inmem', schema=schema, raw_record_column_name=tfma.ARROW_INPUT_COLUMN, telemetry_descriptors=None) tensor_adapter_config = None if schema is not None: tensor_adapter_config = tensor_adapter.TensorAdapterConfig( arrow_schema=coder.ArrowSchema(), tensor_representations=coder.TensorRepresentations()) _ = (pipeline | 'InputFileList' >> beam.Create(input_files) | 'ReadText' >> beam.io.textio.ReadAllFromText() | 'ParseData' >> beam.ParDo( JSONToSerializedExample(eval_column_specs=eval_column_specs, class_list=class_name_list, quantile_list=quantile_list, quantile_index=quantile_index)) | 'ExamplesToRecordBatch' >> coder.BeamSource() | 'ExtractEvaluateAndWriteResults' >> tfma.ExtractEvaluateAndWriteResults( eval_config=tfma_eval_config, writers=me_writers, tensor_adapter_config=tensor_adapter_config))