def testGetModelAndOutputNamesMultiModel(self): eval_config = config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec(name=constants.BASELINE_KEY), config_pb2.ModelSpec(name=constants.CANDIDATE_KEY) ]) self.assertEqual([(constants.BASELINE_KEY, None), (constants.CANDIDATE_KEY, None)], util.StandardExtracts({ constants.PREDICTIONS_KEY: { constants.BASELINE_KEY: np.array([]), constants.CANDIDATE_KEY: np.array([]) } }).get_model_and_output_names(eval_config))
def _data_spec_to_model_spec( self, me_data_spec: me_proto.DataSpec, model_name: Optional[str] = None) -> config_pb2.ModelSpec: """Convert ME DataSpec into TFMA ModelSpec. Args: me_data_spec: Input ME DataSpec. model_name: Input model name. Returns: TFMA ModelSpec. """ if not me_data_spec: return None tfma_model_spec = config_pb2.ModelSpec() if model_name: tfma_model_spec.name = model_name if me_data_spec.HasField('label_key_spec'): tfma_model_spec.label_key = ColumnSpec( me_data_spec.label_key_spec).as_string() + ( constants.Data.K_HOT_KABEL_KEY_SUFFIX if self._class_name_list else '') if me_data_spec.HasField('predicted_score_key_spec'): tfma_model_spec.prediction_key = ColumnSpec( me_data_spec.predicted_score_key_spec).as_string() + ( constants.Data.POINT_KEY_SUFFIX if model_name and model_name.endswith(constants.Data.POINT_KEY_SUFFIX) else '') if me_data_spec.HasField('example_weight_key_spec'): tfma_model_spec.example_weight_key = ColumnSpec( me_data_spec.example_weight_key_spec).as_string() return tfma_model_spec
def testGetModelAndOutputNamesEmptyPredictions(self): eval_config = config_pb2.EvalConfig( model_specs=[config_pb2.ModelSpec()]) self.assertEmpty( util.StandardExtracts({ constants.PREDICTIONS_KEY: {} }).get_model_and_output_names(eval_config))
def testStandardMetricInputsWithCustomLabelKeys(self): example = metric_types.StandardMetricInputs( labels={ 'custom_label': np.array([2]), 'other_label': np.array([0]) }, predictions={'custom_prediction': np.array([0, 0.5, 0.3, 0.9])}, example_weights=np.array([1.0])) eval_config = config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec(label_key='custom_label', prediction_key='custom_prediction') ]) iterator = metric_util.to_label_prediction_example_weight( example, eval_config=eval_config) for expected_label, expected_prediction in zip((0.0, 0.0, 1.0, 0.0), (0.0, 0.5, 0.3, 0.9)): got_label, got_pred, got_example_weight = next(iterator) self.assertAllClose(got_label, np.array([expected_label]), atol=0, rtol=0) self.assertAllClose(got_pred, np.array([expected_prediction]), atol=0, rtol=0) self.assertAllClose(got_example_weight, np.array([1.0]), atol=0, rtol=0)
def testGetFeatureValuesForModelSpecFieldNoValues(self): model_spec = config_pb2.ModelSpec(name='model1', example_weight_key='feature2') extracts = {} got = model_util.get_feature_values_for_model_spec_field( [model_spec], 'example_weight', 'example_weights', extracts) self.assertIsNone(got)
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 testSliceKeys(self, model_names, extracts, slice_specs, expected_slices): eval_config = config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec(name=name) for name in model_names ]) with beam.Pipeline() as pipeline: slice_keys_extracts = ( pipeline | 'CreateTestInput' >> beam.Create(extracts) | 'ExtractSlices' >> slice_key_extractor.ExtractSliceKeys( slice_spec=slice_specs, eval_config=eval_config)) def check_result(got): try: self.assertLen(got, 2) got_results = [] for item in got: self.assertIn(constants.SLICE_KEY_TYPES_KEY, item) got_results.append( sorted(item[constants.SLICE_KEY_TYPES_KEY])) self.assertCountEqual(got_results, expected_slices) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(slice_keys_extracts, check_result)
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 testBatchSizeLimit(self): temp_export_dir = self._getExportDir() _, export_dir = batch_size_limited_classifier.simple_batch_size_limited_classifier( None, temp_export_dir) eval_shared_model = self.createTestEvalSharedModel( eval_saved_model_path=export_dir, tags=[tf.saved_model.SERVING]) eval_config = config_pb2.EvalConfig(model_specs=[config_pb2.ModelSpec()]) schema = text_format.Parse( """ feature { name: "classes" type: BYTES } feature { name: "scores" type: FLOAT } feature { name: "labels" type: BYTES } """, schema_pb2.Schema()) tfx_io = test_util.InMemoryTFExampleRecord( 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()) feature_extractor = features_extractor.FeaturesExtractor( eval_config=eval_config, tensor_representations=tensor_adapter_config.tensor_representations) prediction_extractor = predictions_extractor.PredictionsExtractor( eval_config=eval_config, eval_shared_model=eval_shared_model) examples = [] for _ in range(4): examples.append( self._makeExample(classes='first', scores=0.0, labels='third')) with beam.Pipeline() as pipeline: predict_extracts = ( pipeline | 'Create' >> beam.Create([e.SerializeToString() for e in examples], reshuffle=False) | 'BatchExamples' >> tfx_io.BeamSource(batch_size=1) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | feature_extractor.stage_name >> feature_extractor.ptransform | prediction_extractor.stage_name >> prediction_extractor.ptransform) def check_result(got): try: self.assertLen(got, 4) # We can't verify the actual predictions, but we can verify the keys. for item in got: self.assertIn(constants.PREDICTIONS_KEY, item) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(predict_extracts, check_result, label='result')
def testCustomTFMetricWithPadding(self, example_indices, expected): computation = tf_metric_wrapper.tf_metric_computations( [ _CustomMetric(name='custom_label', update_y_pred=False), _CustomMetric(name='custom_pred', update_y_pred=True), ], eval_config=config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec(padding_options=config_pb2.PaddingOptions( label_int_padding=-1, prediction_float_padding=-1.0, )) ]), example_weighted=True)[0] examples = [{ 'labels': np.array([1], dtype=np.int64), 'predictions': np.array([0.1, 0.2, 0.3, 0.0]), 'example_weights': np.array([1.0]) }, { 'labels': np.array([1, 2], dtype=np.int64), 'predictions': np.array([0.1, 0.2, 0.0]), 'example_weights': np.array([1.0]) }, { 'labels': np.array([1, 2, 3], dtype=np.int64), 'predictions': np.array([0.1, 0.2, 0.3]), 'example_weights': np.array([2.0]) }] with beam.Pipeline() as pipeline: # pylint: disable=no-value-for-parameter result = ( pipeline | 'Create' >> beam.Create([examples[i] for i in example_indices]) | 'Process' >> beam.Map(metric_util.to_standard_metric_inputs) | 'AddSlice' >> beam.Map(lambda x: ((), x)) | 'Combine' >> beam.CombinePerKey(computation.combiner)) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) got_slice_key, got_metrics = got[0] self.assertEqual(got_slice_key, ()) custom_label_key = metric_types.MetricKey( name='custom_label', example_weighted=True) custom_pred_key = metric_types.MetricKey( name='custom_pred', example_weighted=True) self.assertDictElementsAlmostEqual(got_metrics, expected) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
def testGetFeatureValuesForModelSpecFieldNoValues(self): model_spec = config_pb2.ModelSpec(name='model1', example_weight_key='feature2') extracts = { constants.ARROW_RECORD_BATCH_KEY: pa.RecordBatch.from_arrays([pa.array([1])], ['dummy']), } got = model_util.get_feature_values_for_model_spec_field( [model_spec], 'example_weight', 'example_weights', extracts) self.assertIsNone(got)
def testGetModelAndOutputNamesMultiOutput(self): eval_config = config_pb2.EvalConfig( model_specs=[config_pb2.ModelSpec()]) self.assertEqual([(None, 'output1'), (None, 'output2')], util.StandardExtracts({ constants.PREDICTIONS_KEY: { 'output1': np.array([]), 'output2': np.array([]) } }).get_model_and_output_names(eval_config))
def get_evaluation_config( problem_type: constants.ProblemType, evaluation_column_specs: EvaluationColumnSpecs, slice_features: List[List[ColumnSpec]], class_names: Optional[List[Text]] = None, positive_class_names: Optional[List[Text]] = None, top_k_list: Optional[List[int]] = None ) -> model_evaluation_pb2.EvaluationConfig: """Build a Model Evaluation Configuration. Args: problem_type: One of the ProblemType enum. evaluation_column_specs: column specs necessary for parsing evaluation data. slice_features: List of slice specs, each a list of keys to slice. The default slice over all values will automatically be added. class_names: For classification-type problems, a list of string names for classes. positive_class_names: For classification-type problems, a list of string names for classes to be treated as positively valued. top_k_list: For classification-type problems, if specified, a list of top-k aggregations. Returns: An EvaluationConfig. """ tfma_eval_config = config_pb2.EvalConfig() tfma_eval_config.model_specs.append( config_pb2.ModelSpec( prediction_key=evaluation_column_specs.predicted_score_column_spec .as_string(), prediction_keys=None, label_key=evaluation_column_specs.ground_truth_column_spec.as_string( ), label_keys=None)) metric_specs = _get_metric_specs(problem_type, class_names, positive_class_names, top_k_list) assert metric_specs, 'At least one metric_spec must be defined %r' % metric_specs tfma_eval_config.metrics_specs.extend(metric_specs) slicing_specs = _get_tfma_slicing_specs(slice_features) assert slicing_specs, 'At least one slicing_spec must be defined %r' % slicing_specs tfma_eval_config.slicing_specs.extend(slicing_specs) adapter = tfma_adapter.TFMAToME( class_name_list=class_names, predicted_label_column_spec=evaluation_column_specs .predicted_label_column_spec, predicted_label_id_column_spec=evaluation_column_specs .predicted_label_id_column_spec) return adapter.eval_config(tfma_eval_config)
def testBatchedPredict(self): temp_eval_export_dir = self._getEvalExportDir() _, eval_export_dir = linear_classifier.simple_linear_classifier( None, temp_eval_export_dir) eval_shared_model = model_eval_lib.default_eval_shared_model( eval_saved_model_path=eval_export_dir) eval_config = config_pb2.EvalConfig( model_specs=[config_pb2.ModelSpec()]) with beam.Pipeline() as pipeline: examples = [ self._makeExample(age=3.0, language='english', label=1.0), self._makeExample(age=3.0, language='chinese', label=0.0), self._makeExample(age=4.0, language='english', label=1.0), self._makeExample(age=5.0, language='chinese', label=0.0), ] serialized_examples = [e.SerializeToString() for e in examples] tfx_io = raw_tf_record.RawBeamRecordTFXIO( physical_format='inmemory', raw_record_column_name=constants.ARROW_INPUT_COLUMN, telemetry_descriptors=['TFMATest']) extractor = predict_extractor.PredictExtractor( eval_shared_model, eval_config=eval_config) predict_extracts = ( pipeline | 'Create' >> beam.Create(serialized_examples, reshuffle=False) | 'BatchExamples' >> tfx_io.BeamSource(batch_size=2) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | 'Predict' >> extractor.ptransform) def check_result(got): try: self.assertLen(got, 2) for item in got: self.assertIn(constants.FEATURES_KEY, item) for feature in ('language', 'age'): for features_dict in item[constants.FEATURES_KEY]: self.assertIn(feature, features_dict) self.assertIn(constants.LABELS_KEY, item) self.assertIn(constants.PREDICTIONS_KEY, item) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(predict_extracts, check_result, label='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)))
class CalibrationPlotTest(testutil.TensorflowModelAnalysisTest, parameterized.TestCase): def testCalibrationPlot(self): computations = calibration_plot.CalibrationPlot( num_buckets=10).computations(example_weighted=True) histogram = computations[0] plot = computations[1] example1 = { 'labels': np.array([0.0]), 'predictions': np.array([0.2]), 'example_weights': np.array([1.0]) } example2 = { 'labels': np.array([1.0]), 'predictions': np.array([0.8]), 'example_weights': np.array([2.0]) } example3 = { 'labels': np.array([0.0]), 'predictions': np.array([0.5]), 'example_weights': np.array([3.0]) } example4 = { 'labels': np.array([1.0]), 'predictions': np.array([-0.1]), 'example_weights': np.array([4.0]) } example5 = { 'labels': np.array([1.0]), 'predictions': np.array([0.5]), 'example_weights': np.array([5.0]) } example6 = { 'labels': np.array([1.0]), 'predictions': np.array([0.8]), 'example_weights': np.array([6.0]) } example7 = { 'labels': np.array([0.0]), 'predictions': np.array([0.2]), 'example_weights': np.array([7.0]) } example8 = { 'labels': np.array([1.0]), 'predictions': np.array([1.1]), 'example_weights': np.array([8.0]) } with beam.Pipeline() as pipeline: # pylint: disable=no-value-for-parameter result = ( pipeline | 'Create' >> beam.Create([ example1, example2, example3, example4, example5, example6, example7, example8 ]) | 'Process' >> beam.Map(metric_util.to_standard_metric_inputs) | 'AddSlice' >> beam.Map(lambda x: ((), x)) | 'ComputeHistogram' >> beam.CombinePerKey(histogram.combiner) | 'ComputePlot' >> beam.Map(lambda x: (x[0], plot.result(x[1])))) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) got_slice_key, got_plots = got[0] self.assertEqual(got_slice_key, ()) self.assertLen(got_plots, 1) key = metric_types.PlotKey(name='calibration_plot', example_weighted=True) self.assertIn(key, got_plots) got_plot = got_plots[key] self.assertProtoEquals( """ buckets { lower_threshold_inclusive: -inf upper_threshold_exclusive: 0.0 total_weighted_label { value: 4.0 } total_weighted_refined_prediction { value: -0.4 } num_weighted_examples { value: 4.0 } } buckets { lower_threshold_inclusive: 0.0 upper_threshold_exclusive: 0.1 total_weighted_label { } total_weighted_refined_prediction { } num_weighted_examples { } } buckets { lower_threshold_inclusive: 0.1 upper_threshold_exclusive: 0.2 total_weighted_label { } total_weighted_refined_prediction { } num_weighted_examples { } } buckets { lower_threshold_inclusive: 0.2 upper_threshold_exclusive: 0.3 total_weighted_label { } total_weighted_refined_prediction { value: 1.6 } num_weighted_examples { value: 8.0 } } buckets { lower_threshold_inclusive: 0.3 upper_threshold_exclusive: 0.4 total_weighted_label { } total_weighted_refined_prediction { } num_weighted_examples { } } buckets { lower_threshold_inclusive: 0.4 upper_threshold_exclusive: 0.5 total_weighted_label { } total_weighted_refined_prediction { } num_weighted_examples { } } buckets { lower_threshold_inclusive: 0.5 upper_threshold_exclusive: 0.6 total_weighted_label { value: 5.0 } total_weighted_refined_prediction { value: 4.0 } num_weighted_examples { value: 8.0 } } buckets { lower_threshold_inclusive: 0.6 upper_threshold_exclusive: 0.7 total_weighted_label { } total_weighted_refined_prediction { } num_weighted_examples { } } buckets { lower_threshold_inclusive: 0.7 upper_threshold_exclusive: 0.8 total_weighted_label { } total_weighted_refined_prediction { } num_weighted_examples { } } buckets { lower_threshold_inclusive: 0.8 upper_threshold_exclusive: 0.9 total_weighted_label { value: 8.0 } total_weighted_refined_prediction { value: 6.4 } num_weighted_examples { value: 8.0 } } buckets { lower_threshold_inclusive: 0.9 upper_threshold_exclusive: 1.0 total_weighted_label { } total_weighted_refined_prediction { } num_weighted_examples { } } buckets { lower_threshold_inclusive: 1.0 upper_threshold_exclusive: inf total_weighted_label { value: 8.0 } total_weighted_refined_prediction { value: 8.8 } num_weighted_examples { value: 8.0 } } """, got_plot) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result') @parameterized.named_parameters( { 'testcase_name': 'int_single_model', 'eval_config': config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec(name='model1', label_key='label'), ]), 'schema': text_format.Parse( """ feature { name: "label" type: INT int_domain { min: 5 max: 15 } } """, schema_pb2.Schema()), 'model_names': [''], 'output_names': [''], 'expected_left': 5.0, 'expected_range': 10.0, }, { 'testcase_name': 'int_single_model_right_only', 'eval_config': config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec(name='model1', label_key='label'), ]), 'schema': text_format.Parse( """ feature { name: "label" type: INT int_domain { max: 15 } } """, schema_pb2.Schema()), 'model_names': [''], 'output_names': [''], 'expected_left': 0.0, 'expected_range': 15.0, }, { 'testcase_name': 'int_single_model_schema_missing_domain', 'eval_config': config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec(name='model1', label_key='label'), ]), 'schema': text_format.Parse( """ feature { name: "label" type: FLOAT } """, schema_pb2.Schema()), 'model_names': [''], 'output_names': [''], 'expected_left': 0.0, 'expected_range': 1.0, }, { 'testcase_name': 'int_single_model_schema_missing_label', 'eval_config': config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec(name='model1', label_key='label'), ]), 'schema': text_format.Parse( """ feature { name: "other_feature" type: BYTES } """, schema_pb2.Schema()), 'model_names': [''], 'output_names': [''], 'expected_left': 0.0, 'expected_range': 1.0, }, { 'testcase_name': 'float_single_model', 'eval_config': config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec(name='model1', label_key='label'), ]), 'schema': text_format.Parse( """ feature { name: "label" type: FLOAT float_domain { min: 5.0 max: 15.0 } } """, schema_pb2.Schema()), 'model_names': [''], 'output_names': [''], 'expected_left': 5.0, 'expected_range': 10.0 }, { 'testcase_name': 'float_single_model_multiple_outputs', 'eval_config': config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec(name='model1', label_keys={ 'output1': 'label1', 'output2': 'label2' }, signature_name='default'), ]), 'schema': text_format.Parse( """ feature { name: "label2" type: FLOAT float_domain { min: 5.0 max: 15.0 } } """, schema_pb2.Schema()), 'model_names': [''], 'output_names': ['output2'], 'expected_left': 5.0, 'expected_range': 10.0 }, { 'testcase_name': 'float_multiple_models', 'eval_config': config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec(name='model1', label_key='label1'), config_pb2.ModelSpec(name='model2', label_key='label2') ]), 'schema': text_format.Parse( """ feature { name: "label2" type: FLOAT float_domain { min: 5.0 max: 15.0 } } """, schema_pb2.Schema()), 'model_names': ['model2'], 'output_names': [''], 'expected_left': 5.0, 'expected_range': 10.0 }) def testCalibrationPlotWithSchema(self, eval_config, schema, model_names, output_names, expected_left, expected_range): computations = calibration_plot.CalibrationPlot( num_buckets=10).computations(eval_config=eval_config, schema=schema, model_names=model_names, output_names=output_names) histogram = computations[0] self.assertEqual(expected_left, histogram.combiner._left) self.assertEqual(expected_range, histogram.combiner._range)
def testUnbatchExtractor(self): model_spec = config_pb2.ModelSpec(label_key='label', example_weight_key='example_weight') eval_config = config_pb2.EvalConfig(model_specs=[model_spec]) feature_extractor = features_extractor.FeaturesExtractor(eval_config) label_extractor = labels_extractor.LabelsExtractor(eval_config) example_weight_extractor = ( example_weights_extractor.ExampleWeightsExtractor(eval_config)) predict_extractor = predictions_extractor.PredictionsExtractor( eval_config) unbatch_inputs_extractor = unbatch_extractor.UnbatchExtractor() schema = text_format.Parse( """ feature { name: "label" type: FLOAT } 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 = test_util.InMemoryTFExampleRecord( schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN) examples = [ self._makeExample(label=1.0, example_weight=0.5, fixed_int=1, fixed_float=1.0, fixed_string='fixed_string1'), self._makeExample(label=0.0, example_weight=0.0, fixed_int=1, fixed_float=1.0, fixed_string='fixed_string2'), self._makeExample(label=0.0, example_weight=1.0, 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 | label_extractor.stage_name >> label_extractor.ptransform | example_weight_extractor.stage_name >> example_weight_extractor.ptransform | predict_extractor.stage_name >> predict_extractor.ptransform | unbatch_inputs_extractor.stage_name >> unbatch_inputs_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 3) self.assertDictElementsAlmostEqual( got[0][constants.FEATURES_KEY], { 'fixed_int': np.array([1]), 'fixed_float': np.array([1.0]), }) self.assertEqual( got[0][constants.FEATURES_KEY]['fixed_string'], np.array([b'fixed_string1'])) self.assertAlmostEqual(got[0][constants.LABELS_KEY], np.array([1.0])) self.assertAlmostEqual( got[0][constants.EXAMPLE_WEIGHTS_KEY], np.array([0.5])) self.assertDictElementsAlmostEqual( got[1][constants.FEATURES_KEY], { 'fixed_int': np.array([1]), 'fixed_float': np.array([1.0]), }) self.assertEqual( got[1][constants.FEATURES_KEY]['fixed_string'], np.array([b'fixed_string2'])) self.assertAlmostEqual(got[1][constants.LABELS_KEY], np.array([0.0])) self.assertAlmostEqual( got[1][constants.EXAMPLE_WEIGHTS_KEY], np.array([0.0])) self.assertDictElementsAlmostEqual( got[2][constants.FEATURES_KEY], { 'fixed_int': np.array([2]), 'fixed_float': np.array([0.0]), }) self.assertEqual( got[2][constants.FEATURES_KEY]['fixed_string'], np.array([b'fixed_string3'])) self.assertAlmostEqual(got[2][constants.LABELS_KEY], np.array([0.0])) self.assertAlmostEqual( got[2][constants.EXAMPLE_WEIGHTS_KEY], np.array([1.0])) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
def testPredictionsExtractorWithMultiModels(self): temp_export_dir = self._getExportDir() export_dir1, _ = multi_head.simple_multi_head(temp_export_dir, None) export_dir2, _ = multi_head.simple_multi_head(temp_export_dir, None) eval_config = config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec(name='model1'), config_pb2.ModelSpec(name='model2') ]) eval_shared_model1 = self.createTestEvalSharedModel( eval_saved_model_path=export_dir1, tags=[tf.saved_model.SERVING]) eval_shared_model2 = self.createTestEvalSharedModel( eval_saved_model_path=export_dir2, tags=[tf.saved_model.SERVING]) schema = text_format.Parse( """ feature { name: "age" type: FLOAT } feature { name: "langauge" type: BYTES } feature { name: "english_label" type: FLOAT } feature { name: "chinese_label" type: FLOAT } feature { name: "other_label" type: FLOAT } """, schema_pb2.Schema()) tfx_io = test_util.InMemoryTFExampleRecord( 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()) feature_extractor = features_extractor.FeaturesExtractor(eval_config) prediction_extractor = predictions_extractor.PredictionsExtractor( eval_config=eval_config, eval_shared_model={ 'model1': eval_shared_model1, 'model2': eval_shared_model2 }, tensor_adapter_config=tensor_adapter_config) examples = [ self._makeExample(age=1.0, language='english', english_label=1.0, chinese_label=0.0, other_label=0.0), self._makeExample(age=1.0, language='chinese', english_label=0.0, chinese_label=1.0, other_label=0.0), self._makeExample(age=2.0, language='english', english_label=1.0, chinese_label=0.0, other_label=0.0), self._makeExample(age=2.0, language='other', english_label=0.0, chinese_label=1.0, other_label=1.0) ] 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=4) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | feature_extractor.stage_name >> feature_extractor.ptransform | prediction_extractor.stage_name >> prediction_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) for item in got: # We can't verify the actual predictions, but we can verify the keys self.assertIn(constants.PREDICTIONS_KEY, item) for pred in item[constants.PREDICTIONS_KEY]: for model_name in ('model1', 'model2'): self.assertIn(model_name, pred) for output_name in ('chinese_head', 'english_head', 'other_head'): for pred_key in ('logistic', 'probabilities', 'all_classes'): self.assertIn( output_name + '/' + pred_key, pred[model_name]) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
def testGetModelAndOutputNamesEmptyExtracts(self): eval_config = config_pb2.EvalConfig( model_specs=[config_pb2.ModelSpec()]) self.assertEmpty( util.StandardExtracts({}).get_model_and_output_names(eval_config))
def testPredictionsExtractorWithRegressionModel(self): temp_export_dir = self._getExportDir() export_dir, _ = (fixed_prediction_estimator_extra_fields. simple_fixed_prediction_estimator_extra_fields( temp_export_dir, None)) eval_config = config_pb2.EvalConfig( model_specs=[config_pb2.ModelSpec()]) eval_shared_model = self.createTestEvalSharedModel( eval_saved_model_path=export_dir, tags=[tf.saved_model.SERVING]) schema = text_format.Parse( """ feature { name: "prediction" type: FLOAT } feature { name: "label" 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 = test_util.InMemoryTFExampleRecord( 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()) feature_extractor = features_extractor.FeaturesExtractor(eval_config) prediction_extractor = predictions_extractor.PredictionsExtractor( eval_config=eval_config, eval_shared_model=eval_shared_model, tensor_adapter_config=tensor_adapter_config) examples = [ self._makeExample(prediction=0.2, label=1.0, fixed_int=1, fixed_float=1.0, fixed_string='fixed_string1'), self._makeExample(prediction=0.8, label=0.0, fixed_int=1, fixed_float=1.0, fixed_string='fixed_string2'), self._makeExample(prediction=0.5, label=0.0, fixed_int=2, fixed_float=1.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 | prediction_extractor.stage_name >> prediction_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) self.assertIn(constants.PREDICTIONS_KEY, got[0]) expected_preds = [0.2, 0.8, 0.5] self.assertAlmostEqual(got[0][constants.PREDICTIONS_KEY], expected_preds) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
def testPredictionsExtractorWithSequentialKerasModel(self): # Note that the input will be called 'test_input' model = tf.keras.models.Sequential([ tf.keras.layers.Dense(1, activation=tf.nn.sigmoid, input_shape=(2, ), name='test') ]) model.compile(optimizer=tf.keras.optimizers.Adam(lr=.001), loss=tf.keras.losses.binary_crossentropy, metrics=['accuracy']) train_features = {'test_input': [[0.0, 0.0], [1.0, 1.0]]} labels = [[1], [0]] example_weights = [1.0, 0.5] dataset = tf.data.Dataset.from_tensor_slices( (train_features, labels, example_weights)) dataset = dataset.shuffle(buffer_size=1).repeat().batch(2) model.fit(dataset, steps_per_epoch=1) export_dir = self._getExportDir() model.save(export_dir, save_format='tf') eval_config = config_pb2.EvalConfig( model_specs=[config_pb2.ModelSpec()]) eval_shared_model = self.createTestEvalSharedModel( eval_saved_model_path=export_dir, tags=[tf.saved_model.SERVING]) schema = text_format.Parse( """ tensor_representation_group { key: "" value { tensor_representation { key: "test" value { dense_tensor { column_name: "test" shape { dim { size: 2 } } } } } } } feature { name: "test" type: FLOAT } feature { name: "non_model_feature" type: INT } """, schema_pb2.Schema()) tfx_io = test_util.InMemoryTFExampleRecord( 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()) feature_extractor = features_extractor.FeaturesExtractor(eval_config) prediction_extractor = predictions_extractor.PredictionsExtractor( eval_config=eval_config, eval_shared_model=eval_shared_model, tensor_adapter_config=tensor_adapter_config) # Notice that the features are 'test' but the model expects 'test_input'. # This tests that the PredictExtractor properly handles this case. examples = [ self._makeExample( test=[0.0, 0.0], non_model_feature=0), # should be ignored by model self._makeExample( test=[1.0, 1.0], non_model_feature=1), # should be ignored by model ] 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=2) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | feature_extractor.stage_name >> feature_extractor.ptransform | prediction_extractor.stage_name >> prediction_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) # We can't verify the actual predictions, but we can verify the keys. for item in got: self.assertIn(constants.PREDICTIONS_KEY, item) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
def testExampleWeightsExtractorMultiModel(self): model_spec1 = config_pb2.ModelSpec(name='model1', example_weight_key='example_weight') model_spec2 = config_pb2.ModelSpec(name='model2', example_weight_keys={ 'output1': 'example_weight1', 'output2': 'example_weight2' }) eval_config = config_pb2.EvalConfig( model_specs=[model_spec1, model_spec2]) feature_extractor = features_extractor.FeaturesExtractor(eval_config) example_weight_extractor = example_weights_extractor.ExampleWeightsExtractor( eval_config) schema = text_format.Parse( """ feature { name: "example_weight" type: FLOAT } feature { name: "example_weight1" type: FLOAT } feature { name: "example_weight2" type: FLOAT } feature { name: "fixed_int" type: INT } """, schema_pb2.Schema()) tfx_io = test_util.InMemoryTFExampleRecord( schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN) examples = [ self._makeExample(example_weight=0.5, example_weight1=0.5, example_weight2=0.5, fixed_int=1), self._makeExample(example_weight=0.0, example_weight1=0.0, example_weight2=1.0, fixed_int=1) ] 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=2) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | feature_extractor.stage_name >> feature_extractor.ptransform | example_weight_extractor.stage_name >> example_weight_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) for model_name in ('model1', 'model2'): self.assertIn(model_name, got[0][constants.EXAMPLE_WEIGHTS_KEY]) self.assertAllClose( got[0][constants.EXAMPLE_WEIGHTS_KEY]['model1'], np.array([0.5, 0.0])) self.assertAllClose( got[0][constants.EXAMPLE_WEIGHTS_KEY]['model2'], { 'output1': np.array([0.5, 0.0]), 'output2': np.array([0.5, 1.0]) }) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
class ModelUtilTest(testutil.TensorflowModelAnalysisTest, parameterized.TestCase): def createDenseInputsSchema(self): return text_format.Parse( """ tensor_representation_group { key: "" value { tensor_representation { key: "input_1" value { dense_tensor { column_name: "input_1" shape { dim { size: 1 } } } } } tensor_representation { key: "input_2" value { dense_tensor { column_name: "input_2" shape { dim { size: 1 } } } } } } } feature { name: "input_1" type: FLOAT } feature { name: "input_2" type: FLOAT } feature { name: "non_model_feature" type: INT } """, schema_pb2.Schema()) def createModelWithSingleInput(self, save_as_keras): input_layer = tf.keras.layers.Input(shape=(1, ), name='input') output_layer = tf.keras.layers.Dense( 1, activation=tf.nn.sigmoid)(input_layer) model = tf.keras.models.Model(input_layer, output_layer) @tf.function def serving_default(s): return model(s) input_spec = { 'input': tf.TensorSpec(shape=(None, 1), dtype=tf.string, name='input'), } signatures = { 'serving_default': serving_default.get_concrete_function(input_spec), 'custom_signature': serving_default.get_concrete_function(input_spec), } export_path = tempfile.mkdtemp() if save_as_keras: model.save(export_path, save_format='tf', signatures=signatures) else: tf.saved_model.save(model, export_path, signatures=signatures) return export_path def createModelWithMultipleDenseInputs(self, save_as_keras): input1 = tf.keras.layers.Input(shape=(1, ), name='input_1') input2 = tf.keras.layers.Input(shape=(1, ), name='input_2') inputs = [input1, input2] input_layer = tf.keras.layers.concatenate(inputs) output_layer = tf.keras.layers.Dense(1, activation=tf.nn.sigmoid, name='output')(input_layer) model = tf.keras.models.Model(inputs, output_layer) # Add custom attribute to model to test callables stored as attributes model.custom_attribute = tf.keras.models.Model(inputs, output_layer) @tf.function def serving_default(serialized_tf_examples): parsed_features = tf.io.parse_example( serialized_tf_examples, { 'input_1': tf.io.FixedLenFeature([1], dtype=tf.float32), 'input_2': tf.io.FixedLenFeature([1], dtype=tf.float32) }) return model(parsed_features) @tf.function def custom_single_output(features): return model(features) @tf.function def custom_multi_output(features): return {'output1': model(features), 'output2': model(features)} input_spec = tf.TensorSpec(shape=(None, ), dtype=tf.string, name='examples') custom_input_spec = { 'input_1': tf.TensorSpec(shape=(None, 1), dtype=tf.float32, name='input_1'), 'input_2': tf.TensorSpec(shape=(None, 1), dtype=tf.float32, name='input_2') } signatures = { 'serving_default': serving_default.get_concrete_function(input_spec), 'custom_single_output': custom_single_output.get_concrete_function(custom_input_spec), 'custom_multi_output': custom_multi_output.get_concrete_function(custom_input_spec) } export_path = tempfile.mkdtemp() if save_as_keras: model.save(export_path, save_format='tf', signatures=signatures) else: tf.saved_model.save(model, export_path, signatures=signatures) return export_path def createModelWithInvalidOutputShape(self): input1 = tf.keras.layers.Input(shape=(1, ), name='input_1') input2 = tf.keras.layers.Input(shape=(1, ), name='input_2') inputs = [input1, input2] input_layer = tf.keras.layers.concatenate(inputs) output_layer = tf.keras.layers.Dense(2, activation=tf.nn.sigmoid, name='output')(input_layer) # Flatten the layer such that the first dimension no longer corresponds # with the batch size. reshape_layer = tf.keras.layers.Lambda(lambda x: tf.reshape(x, [-1]), name='reshape')(output_layer) model = tf.keras.models.Model(inputs, reshape_layer) @tf.function def serving_default(serialized_tf_examples): parsed_features = tf.io.parse_example( serialized_tf_examples, { 'input_1': tf.io.FixedLenFeature([1], dtype=tf.float32), 'input_2': tf.io.FixedLenFeature([1], dtype=tf.float32) }) return model(parsed_features) input_spec = tf.TensorSpec(shape=(None, ), dtype=tf.string, name='examples') signatures = { 'serving_default': serving_default.get_concrete_function(input_spec), } export_path = tempfile.mkdtemp() model.save(export_path, save_format='tf', signatures=signatures) return export_path def createModelWithMultipleMixedInputs(self, save_as_keras): dense_input = tf.keras.layers.Input(shape=(2, ), name='input_1', dtype=tf.int64) dense_float_input = tf.cast(dense_input, tf.float32) sparse_input = tf.keras.layers.Input(shape=(1, ), name='input_2', sparse=True) dense_sparse_input = tf.keras.layers.Dense( 1, name='dense_input2')(sparse_input) ragged_input = tf.keras.layers.Input(shape=(None, ), name='input_3', ragged=True) dense_ragged_input = tf.keras.layers.Lambda(lambda x: x.to_tensor())( ragged_input) dense_ragged_input.set_shape((None, 1)) inputs = [dense_input, sparse_input, ragged_input] input_layer = tf.keras.layers.concatenate( [dense_float_input, dense_sparse_input, dense_ragged_input]) output_layer = tf.keras.layers.Dense( 1, activation=tf.nn.sigmoid)(input_layer) model = tf.keras.models.Model(inputs, output_layer) @tf.function def serving_default(features): return model(features) input_spec = { 'input_1': tf.TensorSpec(shape=(None, 2), dtype=tf.int64, name='input_1'), 'input_2': tf.SparseTensorSpec(shape=(None, 1), dtype=tf.float32), 'input_3': tf.RaggedTensorSpec(shape=(None, 1), dtype=tf.float32) } signatures = { 'serving_default': serving_default.get_concrete_function(input_spec), 'custom_signature': serving_default.get_concrete_function(input_spec), } export_path = tempfile.mkdtemp() if save_as_keras: model.save(export_path, save_format='tf', signatures=signatures) else: tf.saved_model.save(model, export_path, signatures=signatures) return export_path def testFilterByInputNames(self): tensors = { 'f1': tf.constant([[1.1], [2.1]], dtype=tf.float32), 'f2': tf.constant([[1], [2]], dtype=tf.int64), 'f3': tf.constant([['hello'], ['world']], dtype=tf.string) } filtered_tensors = model_util.filter_by_input_names( tensors, ['f1', 'f3']) self.assertLen(filtered_tensors, 2) self.assertAllEqual(tf.constant([[1.1], [2.1]], dtype=tf.float32), filtered_tensors['f1']) self.assertAllEqual( tf.constant([['hello'], ['world']], dtype=tf.string), filtered_tensors['f3']) @parameterized.named_parameters( ('one_baseline', text_format.Parse( """ model_specs { name: "candidate" } model_specs { name: "baseline" is_baseline: true } """, config_pb2.EvalConfig()), text_format.Parse( """ name: "baseline" is_baseline: true """, config_pb2.ModelSpec())), ('no_baseline', text_format.Parse( """ model_specs { name: "candidate" } """, config_pb2.EvalConfig()), None), ) def test_get_baseline_model(self, eval_config, expected_baseline_model_spec): self.assertEqual(expected_baseline_model_spec, model_util.get_baseline_model_spec(eval_config)) @parameterized.named_parameters( ('one_non_baseline', text_format.Parse( """ model_specs { name: "candidate" } model_specs { name: "baseline" is_baseline: true } """, config_pb2.EvalConfig()), [ text_format.Parse( """ name: "candidate" """, config_pb2.ModelSpec()) ]), ('no_non_baseline', text_format.Parse( """ model_specs { name: "baseline" is_baseline: true } """, config_pb2.EvalConfig()), []), ) def test_get_non_baseline_model(self, eval_config, expected_non_baseline_model_specs): self.assertCountEqual( expected_non_baseline_model_specs, model_util.get_non_baseline_model_specs(eval_config)) def testFilterByInputNamesKeras(self): tensors = { 'f1': tf.constant([[1.1], [2.1]], dtype=tf.float32), 'f2': tf.constant([[1], [2]], dtype=tf.int64), 'f3': tf.constant([['hello'], ['world']], dtype=tf.string) } filtered_tensors = model_util.filter_by_input_names( tensors, [ 'f1' + model_util.KERAS_INPUT_SUFFIX, 'f3' + model_util.KERAS_INPUT_SUFFIX ]) self.assertLen(filtered_tensors, 2) self.assertAllEqual( tf.constant([[1.1], [2.1]], dtype=tf.float32), filtered_tensors['f1' + model_util.KERAS_INPUT_SUFFIX]) self.assertAllEqual( tf.constant([['hello'], ['world']], dtype=tf.string), filtered_tensors['f3' + model_util.KERAS_INPUT_SUFFIX]) @parameterized.named_parameters( ('output_name_and_label_key', config_pb2.ModelSpec(label_key='label'), 'output', 'label'), ('output_name_and_label_keys', config_pb2.ModelSpec(label_keys={'output': 'label'}), 'output', 'label'), ('output_name_and_no_label_keys', config_pb2.ModelSpec(), 'output', None), ('no_output_name_and_label_key', config_pb2.ModelSpec(label_key='label'), '', 'label'), ('no_output_name_and_no_label_keys', config_pb2.ModelSpec(), '', None)) def testGetLabelKey(self, model_spec, output_name, expected_label_key): self.assertEqual(expected_label_key, model_util.get_label_key(model_spec, output_name)) def testGetLabelKeyNoOutputAndLabelKeys(self): with self.assertRaises(ValueError): model_util.get_label_key( config_pb2.ModelSpec(label_keys={'output1': 'label'}), '') @parameterized.named_parameters( { 'testcase_name': 'single_model_single_key', 'model_specs': [config_pb2.ModelSpec(label_key='feature1')], 'field': 'label_key', 'multi_output_field': 'label_keys', 'expected_values': [ [1.0, 1.1, 1.2], ] }, { 'testcase_name': 'single_model_multi_key', 'model_specs': [ config_pb2.ModelSpec(label_keys={ 'output1': 'feature1', 'output2': 'feature2' }) ], 'field': 'label_key', 'multi_output_field': 'label_keys', 'expected_values': [ { 'output1': [1.0, 1.1, 1.2], 'output2': [2.0, 2.1, 2.2] }, ] }, { 'testcase_name': 'multi_model_single_key', 'model_specs': [ config_pb2.ModelSpec(name='model1', example_weight_key='feature2'), config_pb2.ModelSpec(name='model2', example_weight_key='feature3') ], 'field': 'example_weight_key', 'multi_output_field': 'example_weight_keys', 'expected_values': [ { 'model1': [2.0, 2.1, 2.2], 'model2': [3.0, 3.1, 3.2] }, ] }, { 'testcase_name': 'multi_model_multi_key', 'model_specs': [ config_pb2.ModelSpec(name='model1', prediction_keys={ 'output1': 'feature1', 'output2': 'feature2' }), config_pb2.ModelSpec(name='model2', prediction_keys={ 'output1': 'feature1', 'output3': 'feature3' }) ], 'field': 'prediction_key', 'multi_output_field': 'prediction_keys', 'expected_values': [ { 'model1': { 'output1': [1.0, 1.1, 1.2], 'output2': [2.0, 2.1, 2.2] }, 'model2': { 'output1': [1.0, 1.1, 1.2], 'output3': [3.0, 3.1, 3.2] } }, ] }, ) def testGetFeatureValuesForModelSpecField(self, model_specs, field, multi_output_field, expected_values): extracts = { # Only need the num_rows from RecordBatch so use fake array of same len # as features. constants.ARROW_RECORD_BATCH_KEY: pa.RecordBatch.from_arrays([pa.array([1])], ['dummy']), constants.FEATURES_KEY: [ { 'feature1': [1.0, 1.1, 1.2], 'feature2': [2.0, 2.1, 2.2], 'feature3': [3.0, 3.1, 3.2], }, ] } got = model_util.get_feature_values_for_model_spec_field( model_specs, field, multi_output_field, extracts) self.assertAlmostEqual(expected_values, got) @parameterized.named_parameters( { 'testcase_name': 'single_model_single_key', 'model_specs': [config_pb2.ModelSpec(label_key='feature2')], 'field': 'label_key', 'multi_output_field': 'label_keys', 'expected_values': [ [4.0, 4.1, 4.2], ] }, { 'testcase_name': 'single_model_multi_key', 'model_specs': [ config_pb2.ModelSpec(label_keys={ 'output1': 'feature1', 'output2': 'feature2' }) ], 'field': 'label_key', 'multi_output_field': 'label_keys', 'expected_values': [ { 'output1': [1.0, 1.1, 1.2], 'output2': [4.0, 4.1, 4.2] }, ] }, ) def testGetFeatureValuesForModelSpecFieldWithSingleModelTransforedFeatures( self, model_specs, field, multi_output_field, expected_values): extracts = { # Only need the num_rows from RecordBatch so use fake array of same len # as features. constants.ARROW_RECORD_BATCH_KEY: pa.RecordBatch.from_arrays([pa.array([1])], ['dummy']), constants.FEATURES_KEY: [ { 'feature1': [1.0, 1.1, 1.2], 'feature2': [2.0, 2.1, 2.2], }, ], constants.TRANSFORMED_FEATURES_KEY: [ { 'feature2': [4.0, 4.1, 4.2], }, ] } got = model_util.get_feature_values_for_model_spec_field( model_specs, field, multi_output_field, extracts) self.assertAlmostEqual(expected_values, got) @parameterized.named_parameters( { 'testcase_name': 'multi_model_single_key', 'model_specs': [ config_pb2.ModelSpec(name='model1', example_weight_key='feature2'), config_pb2.ModelSpec(name='model2', example_weight_key='feature3') ], 'field': 'example_weight_key', 'multi_output_field': 'example_weight_keys', 'expected_values': [ { 'model1': [4.0, 4.1, 4.2], 'model2': [7.0, 7.1, 7.2] }, ] }, { 'testcase_name': 'multi_model_multi_key', 'model_specs': [ config_pb2.ModelSpec(name='model1', example_weight_keys={ 'output1': 'feature1', 'output2': 'feature2' }), config_pb2.ModelSpec(name='model2', example_weight_keys={ 'output1': 'feature1', 'output3': 'feature3' }) ], 'field': 'example_weight_key', 'multi_output_field': 'example_weight_keys', 'expected_values': [ { 'model1': { 'output1': [1.0, 1.1, 1.2], 'output2': [4.0, 4.1, 4.2] }, 'model2': { 'output1': [1.0, 1.1, 1.2], 'output3': [7.0, 7.1, 7.2] } }, ] }, ) def testGetFeatureValuesForModelSpecFieldWithMultiModelTransforedFeatures( self, model_specs, field, multi_output_field, expected_values): extracts = { # Only need the num_rows from RecordBatch so use fake array of same len # as features. constants.ARROW_RECORD_BATCH_KEY: pa.RecordBatch.from_arrays([pa.array([1])], ['dummy']), constants.FEATURES_KEY: [ { 'feature1': [1.0, 1.1, 1.2], 'feature2': [2.0, 2.1, 2.2], }, ], constants.TRANSFORMED_FEATURES_KEY: [ { 'model1': { 'feature2': [4.0, 4.1, 4.2], 'feature3': [5.0, 5.1, 5.2] }, 'model2': { 'feature2': [6.0, 6.1, 6.2], 'feature3': [7.0, 7.1, 7.2] } }, ] } got = model_util.get_feature_values_for_model_spec_field( model_specs, field, multi_output_field, extracts) self.assertAlmostEqual(expected_values, got) def testGetFeatureValuesForModelSpecFieldNoValues(self): model_spec = config_pb2.ModelSpec(name='model1', example_weight_key='feature2') extracts = { constants.ARROW_RECORD_BATCH_KEY: pa.RecordBatch.from_arrays([pa.array([1])], ['dummy']), } got = model_util.get_feature_values_for_model_spec_field( [model_spec], 'example_weight', 'example_weights', extracts) self.assertIsNone(got) @parameterized.named_parameters( ('keras_serving_default', True, 'serving_default'), ('keras_custom_signature', True, 'custom_signature'), ('tf2_serving_default', False, 'serving_default'), ('tf2_custom_signature', False, 'custom_signature')) def testGetCallableWithSignatures(self, save_as_keras, signature_name): export_path = self.createModelWithSingleInput(save_as_keras) if save_as_keras: model = tf.keras.models.load_model(export_path) else: model = tf.compat.v1.saved_model.load_v2(export_path) self.assertIsNotNone(model_util.get_callable(model, signature_name)) @parameterized.named_parameters(('keras', True), ('tf2', False)) def testGetCallableWithMissingSignatures(self, save_as_keras): export_path = self.createModelWithSingleInput(save_as_keras) if save_as_keras: model = tf.keras.models.load_model(export_path) else: model = tf.compat.v1.saved_model.load_v2(export_path) with self.assertRaises(ValueError): model_util.get_callable(model, 'non_existent') @unittest.skipIf(_TF_MAJOR_VERSION < 2, 'not all input types supported for TF1') def testGetCallableWithKerasModel(self): export_path = self.createModelWithMultipleMixedInputs(True) model = tf.keras.models.load_model(export_path) self.assertEqual(model, model_util.get_callable(model)) @parameterized.named_parameters( ('keras_serving_default', True, 'serving_default'), ('keras_custom_signature', True, 'custom_signature'), ('tf2_serving_default', False, None), ('tf2_custom_signature', False, 'custom_signature')) def testGetInputSpecsWithSignatures(self, save_as_keras, signature_name): export_path = self.createModelWithSingleInput(save_as_keras) if save_as_keras: model = tf.keras.models.load_model(export_path) else: model = tf.compat.v1.saved_model.load_v2(export_path) self.assertEqual( { 'input': tf.TensorSpec(name='input', shape=(None, 1), dtype=tf.string), }, model_util.get_input_specs(model, signature_name)) @parameterized.named_parameters(('keras', True), ('tf2', False)) def testGetInputSpecsWithMissingSignatures(self, save_as_keras): export_path = self.createModelWithSingleInput(save_as_keras) if save_as_keras: model = tf.keras.models.load_model(export_path) else: model = tf.compat.v1.saved_model.load_v2(export_path) with self.assertRaises(ValueError): model_util.get_callable(model, 'non_existent') @unittest.skipIf(_TF_MAJOR_VERSION < 2, 'not all input types supported for TF1') def testGetInputSpecsWithKerasModel(self): export_path = self.createModelWithMultipleMixedInputs(True) model = tf.keras.models.load_model(export_path) # Some versions of TF set the TensorSpec.name and others do not. Since we # don't care about the name, clear it from the output for testing purposes specs = model_util.get_input_specs(model) for k, v in specs.items(): if isinstance(v, tf.TensorSpec): specs[k] = tf.TensorSpec(shape=v.shape, dtype=v.dtype) self.assertEqual( { 'input_1': tf.TensorSpec(shape=(None, 2), dtype=tf.int64), 'input_2': tf.SparseTensorSpec(shape=(None, 1), dtype=tf.float32), 'input_3': tf.RaggedTensorSpec(shape=(None, None), dtype=tf.float32), }, specs) def testInputSpecsToTensorRepresentations(self): tensor_representations = model_util.input_specs_to_tensor_representations( { 'input_1': tf.TensorSpec(shape=(None, 2), dtype=tf.int64), 'input_2': tf.SparseTensorSpec(shape=(None, 1), dtype=tf.float32), 'input_3': tf.RaggedTensorSpec(shape=(None, None), dtype=tf.float32), }) dense_tensor_representation = text_format.Parse( """ dense_tensor { column_name: "input_1" shape { dim { size: 2 } } } """, schema_pb2.TensorRepresentation()) sparse_tensor_representation = text_format.Parse( """ varlen_sparse_tensor { column_name: "input_2" } """, schema_pb2.TensorRepresentation()) ragged_tensor_representation = text_format.Parse( """ ragged_tensor { feature_path { step: "input_3" } } """, schema_pb2.TensorRepresentation()) self.assertEqual( { 'input_1': dense_tensor_representation, 'input_2': sparse_tensor_representation, 'input_3': ragged_tensor_representation }, tensor_representations) def testInputSpecsToTensorRepresentationsRaisesWithUnknownDims(self): with self.assertRaises(ValueError): model_util.input_specs_to_tensor_representations({ 'input_1': tf.TensorSpec(shape=(None, None), dtype=tf.int64), }) @parameterized.named_parameters( ('keras_default', True, { constants.PREDICTIONS_KEY: { '': [None] } }, None, False, True, 1), ('tf_default', False, { constants.PREDICTIONS_KEY: { '': [None] } }, None, False, True, 1), ('keras_serving_default', True, { constants.PREDICTIONS_KEY: { '': ['serving_default'] } }, None, False, True, 1), ('tf_serving_default', False, { constants.PREDICTIONS_KEY: { '': ['serving_default'] } }, None, False, True, 1), ('keras_custom_single_output', True, { constants.PREDICTIONS_KEY: { '': ['custom_single_output'] } }, None, False, True, 1), ('tf_custom_single_output', False, { constants.PREDICTIONS_KEY: { '': ['custom_single_output'] } }, None, False, True, 1), ('keras_custom_multi_output', True, { constants.PREDICTIONS_KEY: { '': ['custom_multi_output'] } }, None, False, True, 2), ('tf_custom_multi_output', False, { constants.PREDICTIONS_KEY: { '': ['custom_multi_output'] } }, None, False, True, 2), ('multi_model', True, { constants.PREDICTIONS_KEY: { 'model1': ['custom_multi_output'], 'model2': ['custom_multi_output'] } }, None, False, True, 2), ('default_signatures', True, { constants.PREDICTIONS_KEY: { '': [], } }, ['unknown', 'custom_single_output'], False, True, 1), ('keras_prefer_dict_outputs', True, { constants.FEATURES_KEY: { '': [], } }, ['unknown', 'custom_single_output', 'custom_multi_output' ], True, True, 3), ('tf_prefer_dict_outputs', False, { constants.FEATURES_KEY: { '': [], } }, ['unknown', 'custom_single_output', 'custom_multi_output' ], True, True, 3), ('custom_attribute', True, { constants.FEATURES_KEY: { '': ['custom_attribute'], } }, None, True, True, 1), ('keras_no_schema', True, { constants.PREDICTIONS_KEY: { '': [None] } }, None, False, False, 1), ('tf_no_schema', False, { constants.PREDICTIONS_KEY: { '': [None] } }, None, False, False, 1), ('preprocessing_function', True, { constants.TRANSFORMED_FEATURES_KEY: { '': ['_plus_one@input_1'] } }, None, False, False, 1), ) @unittest.skipIf(_TF_MAJOR_VERSION < 2, 'not all signatures supported for TF1') 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) tensor_adapter_config = None if use_schema: 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 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, tensor_adapter_config=tensor_adapter_config))) # 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: for entry in got[0][key]: self.assertIsInstance(entry, dict) self.assertLen(entry, expected_num_outputs) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='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 testHasRubberStamp(self): # Model agnostic. self.assertFalse(model_util.has_rubber_stamp(None)) # All non baseline models has rubber stamp. baseline = self.createTestEvalSharedModel( model_name=constants.BASELINE_KEY, is_baseline=True) candidate = self.createTestEvalSharedModel( model_name=constants.CANDIDATE_KEY, rubber_stamp=True) self.assertTrue(model_util.has_rubber_stamp([baseline, candidate])) # Not all non baseline has rubber stamp. candidate_nr = self.createTestEvalSharedModel( model_name=constants.CANDIDATE_KEY) self.assertFalse(model_util.has_rubber_stamp([candidate_nr])) self.assertFalse( model_util.has_rubber_stamp([baseline, candidate, candidate_nr]))
def testGetLabelKeyNoOutputAndLabelKeys(self): with self.assertRaises(ValueError): model_util.get_label_key( config_pb2.ModelSpec(label_keys={'output1': 'label'}), '')
def testTFJSPredictExtractorWithKerasModel(self, multi_model, multi_output): if not _TFJS_IMPORTED: self.skipTest('This test requires TensorFlow JS.') input1 = tf.keras.layers.Input(shape=(1, ), name='input1') input2 = tf.keras.layers.Input(shape=(1, ), name='input2') inputs = [input1, input2] input_layer = tf.keras.layers.concatenate(inputs) output_layers = {} output_layers['output1'] = (tf.keras.layers.Dense( 1, activation=tf.nn.sigmoid, name='output1')(input_layer)) if multi_output: output_layers['output2'] = (tf.keras.layers.Dense( 1, activation=tf.nn.sigmoid, name='output2')(input_layer)) model = tf.keras.models.Model(inputs, output_layers) model.compile(optimizer=tf.keras.optimizers.Adam(lr=.001), loss=tf.keras.losses.binary_crossentropy, metrics=['accuracy']) train_features = {'input1': [[0.0], [1.0]], 'input2': [[1.0], [0.0]]} labels = {'output1': [[1], [0]]} if multi_output: labels['output2'] = [[1], [0]] example_weights = {'output1': [1.0, 0.5]} if multi_output: example_weights['output2'] = [1.0, 0.5] dataset = tf.data.Dataset.from_tensor_slices( (train_features, labels, example_weights)) dataset = dataset.shuffle(buffer_size=1).repeat().batch(2) model.fit(dataset, steps_per_epoch=1) src_model_path = tempfile.mkdtemp() model.save(src_model_path) dst_model_path = tempfile.mkdtemp() converter.convert([ '--input_format=tf_saved_model', '--saved_model_tags=serve', '--signature_name=serving_default', src_model_path, dst_model_path, ]) model_specs = [config_pb2.ModelSpec(name='model1', model_type='tf_js')] if multi_model: model_specs.append( config_pb2.ModelSpec(name='model2', model_type='tf_js')) eval_config = config_pb2.EvalConfig(model_specs=model_specs) eval_shared_models = [ self.createTestEvalSharedModel( model_name='model1', eval_saved_model_path=dst_model_path, model_type='tf_js') ] if multi_model: eval_shared_models.append( self.createTestEvalSharedModel( model_name='model2', eval_saved_model_path=dst_model_path, model_type='tf_js')) schema = text_format.Parse( """ feature { name: "input1" type: FLOAT } feature { name: "input2" type: FLOAT } feature { name: "non_model_feature" type: INT } """, schema_pb2.Schema()) tfx_io = test_util.InMemoryTFExampleRecord( schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN) feature_extractor = features_extractor.FeaturesExtractor(eval_config) predictor = tfjs_predict_extractor.TFJSPredictExtractor( eval_config=eval_config, eval_shared_model=eval_shared_models) examples = [ self._makeExample(input1=0.0, input2=1.0, non_model_feature=0), self._makeExample(input1=1.0, input2=0.0, non_model_feature=1), ] 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=2) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | feature_extractor.stage_name >> feature_extractor.ptransform | predictor.stage_name >> predictor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) got = got[0] self.assertIn(constants.PREDICTIONS_KEY, got) self.assertLen(got[constants.PREDICTIONS_KEY], 2) for item in got[constants.PREDICTIONS_KEY]: if multi_model: self.assertIn('model1', item) self.assertIn('model2', item) if multi_output: self.assertIn('Identity', item['model1']) self.assertIn('Identity_1', item['model1']) elif multi_output: self.assertIn('Identity', item) self.assertIn('Identity_1', item) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
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 testPredictionsExtractorWithoutEvalSharedModel(self): model_spec1 = config_pb2.ModelSpec(name='model1', prediction_key='prediction') model_spec2 = config_pb2.ModelSpec(name='model2', prediction_keys={ 'output1': 'prediction1', 'output2': 'prediction2' }) eval_config = config_pb2.EvalConfig( model_specs=[model_spec1, model_spec2]) feature_extractor = features_extractor.FeaturesExtractor(eval_config) prediction_extractor = predictions_extractor.PredictionsExtractor( eval_config) schema = text_format.Parse( """ feature { name: "prediction" type: FLOAT } feature { name: "prediction1" type: FLOAT } feature { name: "prediction2" type: FLOAT } feature { name: "fixed_int" type: INT } """, schema_pb2.Schema()) tfx_io = test_util.InMemoryTFExampleRecord( schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN) examples = [ self._makeExample(prediction=1.0, prediction1=1.0, prediction2=0.0, fixed_int=1), self._makeExample(prediction=1.0, prediction1=1.0, prediction2=1.0, fixed_int=1) ] 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=2) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | feature_extractor.stage_name >> feature_extractor.ptransform | prediction_extractor.stage_name >> prediction_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) for model_name in ('model1', 'model2'): self.assertIn(model_name, got[0][constants.PREDICTIONS_KEY][0]) self.assertAlmostEqual( got[0][constants.PREDICTIONS_KEY][0]['model1'], np.array([1.0])) self.assertDictElementsAlmostEqual( got[0][constants.PREDICTIONS_KEY][0]['model2'], { 'output1': np.array([1.0]), 'output2': np.array([0.0]) }) for model_name in ('model1', 'model2'): self.assertIn(model_name, got[0][constants.PREDICTIONS_KEY][1]) self.assertAlmostEqual( got[0][constants.PREDICTIONS_KEY][1]['model1'], np.array([1.0])) self.assertDictElementsAlmostEqual( got[0][constants.PREDICTIONS_KEY][1]['model2'], { 'output1': np.array([1.0]), 'output2': np.array([1.0]) }) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
def testPreprocessedFeaturesExtractor(self, save_as_keras, preprocessing_function_names, expected_extract_keys): export_path = self.createModelWithMultipleDenseInputs(save_as_keras) eval_config = config_pb2.EvalConfig(model_specs=[ config_pb2.ModelSpec( preprocessing_function_names=preprocessing_function_names) ]) eval_shared_model = self.createTestEvalSharedModel( eval_saved_model_path=export_path, tags=[tf.saved_model.SERVING]) schema = self.createDenseInputsSchema() tfx_io = test_util.InMemoryTFExampleRecord( 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()) feature_extractor = features_extractor.FeaturesExtractor(eval_config) transformation_extractor = ( transformed_features_extractor.TransformedFeaturesExtractor( eval_config=eval_config, eval_shared_model=eval_shared_model, tensor_adapter_config=tensor_adapter_config)) 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], reshuffle=False) | 'BatchExamples' >> tfx_io.BeamSource(batch_size=2) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | feature_extractor.stage_name >> feature_extractor.ptransform | transformation_extractor.stage_name >> transformation_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 2) for item in got: for extracts_key, feature_keys in expected_extract_keys.items( ): self.assertIn(extracts_key, item) for value in item[extracts_key]: self.assertEqual(set(feature_keys), set(value.keys()), msg='got={}'.format(item)) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
def testExampleWeightsExtractor(self, example_weight): model_spec = config_pb2.ModelSpec(example_weight_key=example_weight) eval_config = config_pb2.EvalConfig(model_specs=[model_spec]) feature_extractor = features_extractor.FeaturesExtractor(eval_config) example_weight_extractor = ( example_weights_extractor.ExampleWeightsExtractor(eval_config)) example_weight_feature = '' if example_weight is not None: example_weight_feature = """ feature { name: "%s" type: FLOAT } """ % example_weight schema = text_format.Parse( example_weight_feature + """ feature { name: "fixed_int" type: INT } """, schema_pb2.Schema()) tfx_io = test_util.InMemoryTFExampleRecord( schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN) def maybe_add_key(d, key, value): if key is not None: d[key] = value return d example_kwargs = [ maybe_add_key({ 'fixed_int': 1, }, example_weight, 0.5), maybe_add_key({ 'fixed_int': 1, }, example_weight, 0.0), maybe_add_key({ 'fixed_int': 2, }, example_weight, 1.0), ] 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) | 'BatchExamples' >> tfx_io.BeamSource(batch_size=3) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | feature_extractor.stage_name >> feature_extractor.ptransform | example_weight_extractor.stage_name >> example_weight_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) if example_weight: self.assertAllClose( got[0][constants.EXAMPLE_WEIGHTS_KEY], np.array([0.5, 0.0, 1.0])) else: self.assertNotIn(constants.EXAMPLE_WEIGHTS_KEY, got[0]) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
def testBatchSizeLimitWithKerasModel(self): input1 = tf.keras.layers.Input(shape=(1, ), batch_size=1, name='input1') input2 = tf.keras.layers.Input(shape=(1, ), batch_size=1, name='input2') inputs = [input1, input2] input_layer = tf.keras.layers.concatenate(inputs) def add_1(tensor): return tf.add_n([tensor, tf.constant(1.0, shape=(1, 2))]) assert_layer = tf.keras.layers.Lambda(add_1)(input_layer) model = tf.keras.models.Model(inputs, assert_layer) model.compile(optimizer=tf.keras.optimizers.Adam(lr=.001), loss=tf.keras.losses.binary_crossentropy, metrics=['accuracy']) export_dir = self._getExportDir() model.save(export_dir, save_format='tf') eval_config = config_pb2.EvalConfig( model_specs=[config_pb2.ModelSpec()]) eval_shared_model = self.createTestEvalSharedModel( eval_saved_model_path=export_dir, tags=[tf.saved_model.SERVING]) schema = text_format.Parse( """ tensor_representation_group { key: "" value { tensor_representation { key: "input1" value { dense_tensor { column_name: "input1" shape { dim { size: 1 } } } } } tensor_representation { key: "input2" value { dense_tensor { column_name: "input2" shape { dim { size: 1 } } } } } } } feature { name: "input1" type: FLOAT } feature { name: "input2" type: FLOAT } """, schema_pb2.Schema()) tfx_io = test_util.InMemoryTFExampleRecord( 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()) feature_extractor = features_extractor.FeaturesExtractor(eval_config) prediction_extractor = predictions_extractor.PredictionsExtractor( eval_config=eval_config, eval_shared_model=eval_shared_model, tensor_adapter_config=tensor_adapter_config) examples = [] for _ in range(4): examples.append(self._makeExample(input1=0.0, input2=1.0)) with beam.Pipeline() as pipeline: predict_extracts = ( pipeline | 'Create' >> beam.Create( [e.SerializeToString() for e in examples], reshuffle=False) | 'BatchExamples' >> tfx_io.BeamSource(batch_size=1) | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() | feature_extractor.stage_name >> feature_extractor.ptransform | prediction_extractor.stage_name >> prediction_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 4) # We can't verify the actual predictions, but we can verify the keys. for item in got: self.assertIn(constants.PREDICTIONS_KEY, item) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(predict_extracts, check_result, label='result')