def testMakeSklearnPredictExtractor(self): """Tests that predictions are made from extracts for a single model.""" feature_extractor = features_extractor.FeaturesExtractor( self._eval_config) prediction_extractor = ( sklearn_predict_extractor._make_sklearn_predict_extractor( self._eval_shared_model)) with beam.Pipeline() as pipeline: predict_extracts = ( pipeline | 'Create' >> beam.Create( [e.SerializeToString() for e in self._examples]) | 'BatchExamples' >> self._tfx_io.BeamSource() | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() # pylint: disable=no-value-for-parameter | feature_extractor.stage_name >> feature_extractor.ptransform | prediction_extractor.stage_name >> prediction_extractor.ptransform) def check_result(actual): try: for item in actual: self.assertEqual(item['labels'].shape, item['predictions'].shape) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(predict_extracts, check_result)
def 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 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.EvalConfig(model_specs=[config.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) 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(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 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 testMakeSklearnPredictExtractorWithMultiModels(self): """Tests that predictions are made from extracts for multiple models.""" eval_config = config.EvalConfig(model_specs=[ config.ModelSpec(name='model1'), config.ModelSpec(name='model2'), ]) eval_export_dir_1 = os.path.join(self._eval_export_dir, '1') self._create_sklearn_model(eval_export_dir_1) eval_shared_model_1 = sklearn_predict_extractor.custom_eval_shared_model( eval_saved_model_path=eval_export_dir_1, model_name='model1', eval_config=eval_config) eval_export_dir_2 = os.path.join(self._eval_export_dir, '2') self._create_sklearn_model(eval_export_dir_2) eval_shared_model_2 = sklearn_predict_extractor.custom_eval_shared_model( eval_saved_model_path=eval_export_dir_2, model_name='model2', eval_config=eval_config) feature_extractor = features_extractor.FeaturesExtractor( self._eval_config) prediction_extractor = ( sklearn_predict_extractor._make_sklearn_predict_extractor( eval_shared_model={ 'model1': eval_shared_model_1, 'model2': eval_shared_model_2, })) with beam.Pipeline() as pipeline: predict_extracts = ( pipeline | 'Create' >> beam.Create( [e.SerializeToString() for e in self._examples]) | 'BatchExamples' >> self._tfx_io.BeamSource() | 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts() # pylint: disable=no-value-for-parameter | feature_extractor.stage_name >> feature_extractor.ptransform | prediction_extractor.stage_name >> prediction_extractor.ptransform) def check_result(actual): try: for item in actual: self.assertEqual(item['labels'].shape, item['predictions'].shape) self.assertIn('model1', item['predictions'][0]) self.assertIn('model2', item['predictions'][0]) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(predict_extracts, check_result)
def benchmarkMiniPipeline(self): """Benchmark a "mini" TFMA - predict, slice and compute metrics. Runs a "mini" version of TFMA in a Beam pipeline. Records the wall time taken for the whole pipeline. """ self._init_model() pipeline = self._create_beam_pipeline() tfx_io = test_util.InMemoryTFExampleRecord( schema=benchmark_utils.read_schema( self._dataset.tf_metadata_schema_path()), raw_record_column_name=constants.ARROW_INPUT_COLUMN) raw_data = ( pipeline | "Examples" >> beam.Create( self._dataset.read_raw_dataset( deserialize=False, limit=MAX_NUM_EXAMPLES)) | "BatchExamples" >> tfx_io.BeamSource() | "InputsToExtracts" >> tfma.BatchedInputsToExtracts()) _ = ( raw_data | "FeaturesExtractor" >> features_extractor.FeaturesExtractor( eval_config=self._eval_config).ptransform | "LabelsExtractor" >> labels_extractor.LabelsExtractor( eval_config=self._eval_config).ptransform | "ExampleWeightsExtractor" >> example_weights_extractor .ExampleWeightsExtractor(eval_config=self._eval_config).ptransform | "PredictionsExtractor" >> predictions_extractor.PredictionsExtractor( eval_config=self._eval_config, eval_shared_model=self._eval_shared_model).ptransform | "UnbatchExtractor" >> unbatch_extractor.UnbatchExtractor().ptransform | "SliceKeyExtractor" >> tfma.extractors.SliceKeyExtractor().ptransform | "ComputeMetricsPlotsAndValidations" >> metrics_plots_and_validations_evaluator .MetricsPlotsAndValidationsEvaluator( eval_config=self._eval_config, eval_shared_model=self._eval_shared_model).ptransform) start = time.time() result = pipeline.run() result.wait_until_finish() end = time.time() delta = end - start self.report_benchmark( iters=1, wall_time=delta, extras={ "num_examples": self._dataset.num_examples(limit=MAX_NUM_EXAMPLES) })
def testPreprocessedFeaturesExtractor(self, save_as_keras, preprocessing_function_names, expected_extract_keys): export_path = self.createModelWithMultipleDenseInputs(save_as_keras) eval_config = config.EvalConfig(model_specs=[ config.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 testTFlitePredictExtractorWithKerasModel(self, multi_model, multi_output): 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) converter = tf.compat.v2.lite.TFLiteConverter.from_keras_model(model) tflite_model = converter.convert() tflite_model_dir = tempfile.mkdtemp() with tf.io.gfile.GFile(os.path.join(tflite_model_dir, 'tflite'), 'wb') as f: f.write(tflite_model) model_specs = [config.ModelSpec(name='model1', model_type='tf_lite')] if multi_model: model_specs.append( config.ModelSpec(name='model2', model_type='tf_lite')) eval_config = config.EvalConfig(model_specs=model_specs) eval_shared_models = [ self.createTestEvalSharedModel( model_name='model1', eval_saved_model_path=tflite_model_dir, model_type='tf_lite') ] if multi_model: eval_shared_models.append( self.createTestEvalSharedModel( model_name='model2', eval_saved_model_path=tflite_model_dir, model_type='tf_lite')) 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 = tflite_predict_extractor.TFLitePredictExtractor( 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 testPredictionsExtractorWithMultiClassModel(self): temp_export_dir = self._getExportDir() export_dir, _ = dnn_classifier.simple_dnn_classifier( temp_export_dir, None, n_classes=3) 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: "age" type: FLOAT } feature { name: "langauge" type: BYTES } feature { name: "label" 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=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 = [ self._makeExample(age=1.0, language='english', label=0), self._makeExample(age=2.0, language='chinese', label=1), self._makeExample(age=3.0, language='english', label=2), self._makeExample(age=4.0, language='chinese', label=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=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) # We can't verify the actual predictions, but we can verify the keys. for item in got: self.assertIn(constants.PREDICTIONS_KEY, item) for pred_key in ('probabilities', 'all_classes'): self.assertIn(pred_key, item[constants.PREDICTIONS_KEY]) 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.EvalConfig(model_specs=[config.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')
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.EvalConfig(model_specs=[config.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 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.EvalConfig(model_specs=[config.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 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.EvalConfig(model_specs=[ config.ModelSpec(name='model1'), config.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 testSqlSliceKeyExtractorWithMultipleSchema(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 """) ]) feature_extractor = features_extractor.FeaturesExtractor( eval_config=eval_config) slice_key_extractor = sql_slice_key_extractor.SqlSliceKeyExtractor( eval_config) record_batch_1 = pa.RecordBatch.from_arrays([ pa.array([[1], [1], [2]], type=pa.list_(pa.int64())), pa.array([[1.0], [1.0], [2.0]], type=pa.list_(pa.float64())), pa.array([['fixed_string1'], ['fixed_string2'], ['fixed_string3']], type=pa.list_(pa.string())), ], ['fixed_int', 'fixed_float', 'fixed_string']) record_batch_2 = pa.RecordBatch.from_arrays([ pa.array([[1], [1], [2]], type=pa.list_(pa.int64())), pa.array([[1.0], [1.0], [2.0]], type=pa.list_(pa.float64())), pa.array([['fixed_string1'], ['fixed_string2'], ['fixed_string3']], type=pa.list_(pa.string())), pa.array([['extra_field1'], ['extra_field2'], ['extra_field3']], type=pa.list_(pa.string())), ], ['fixed_int', 'fixed_float', 'fixed_string', 'extra_field']) with beam.Pipeline() as pipeline: # pylint: disable=no-value-for-parameter result = ( pipeline | 'Create' >> beam.Create([record_batch_1, record_batch_2], reshuffle=False) | '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, 2) np.testing.assert_equal( got[0][constants.SLICE_KEY_TYPES_KEY], types.VarLenTensorValue.from_dense_rows([ slicer_lib.slice_keys_to_numpy_array([ (('fixed_string', 'fixed_string1'), ) ]), slicer_lib.slice_keys_to_numpy_array([ (('fixed_string', 'fixed_string2'), ) ]), np.array([]) ])) np.testing.assert_equal( got[1][constants.SLICE_KEY_TYPES_KEY], types.VarLenTensorValue.from_dense_rows([ slicer_lib.slice_keys_to_numpy_array([ (('fixed_string', 'fixed_string1'), ) ]), slicer_lib.slice_keys_to_numpy_array([ (('fixed_string', 'fixed_string2'), ) ]), np.array([]) ])) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result)
def testLabelsExtractorMultiModel(self): model_spec1 = config.ModelSpec(name='model1', label_key='label') model_spec2 = config.ModelSpec(name='model2', label_keys={ 'output1': 'label1', 'output2': 'label2' }) eval_config = config.EvalConfig(model_specs=[model_spec1, model_spec2]) feature_extractor = features_extractor.FeaturesExtractor(eval_config) label_extractor = labels_extractor.LabelsExtractor(eval_config) schema = text_format.Parse( """ feature { name: "label" type: FLOAT } feature { name: "label1" type: FLOAT } feature { name: "label2" 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(label=1.0, label1=1.0, label2=0.0, fixed_int=1), self._makeExample(label=1.0, label1=1.0, label2=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 | label_extractor.stage_name >> label_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.LABELS_KEY][0]) self.assertAlmostEqual( got[0][constants.LABELS_KEY][0]['model1'], np.array([1.0])) self.assertDictElementsAlmostEqual( got[0][constants.LABELS_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.LABELS_KEY][1]) self.assertAlmostEqual( got[0][constants.LABELS_KEY][1]['model1'], np.array([1.0])) self.assertDictElementsAlmostEqual( got[0][constants.LABELS_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 testExampleWeightsExtractorMultiOutput(self): model_spec = config_pb2.ModelSpec( example_weight_keys={ 'output1': 'example_weight1', 'output2': 'example_weight2', 'output3': 'example_weight3', }) 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) schema = text_format.Parse( """ 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_weight1=0.5, example_weight2=0.5, fixed_int=1), self._makeExample(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) self.assertAllClose( got[0][constants.EXAMPLE_WEIGHTS_KEY], { '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')
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 testFeaturesExtractor(self): model_spec = config.ModelSpec() eval_config = config.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 = test_util.InMemoryTFExampleRecord( schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN) 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) | 'BatchExamples' >> 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.assertDictElementsAlmostEqual( got[0][constants.FEATURES_KEY][0], { 'fixed_int': np.array([1]), 'fixed_float': np.array([1.0]), }) self.assertEqual( got[0][constants.FEATURES_KEY][0]['fixed_string'], np.array([b'fixed_string1'])) self.assertDictElementsAlmostEqual( got[0][constants.FEATURES_KEY][1], { 'fixed_int': np.array([1]), 'fixed_float': np.array([1.0]), }) self.assertEqual( got[0][constants.FEATURES_KEY][1]['fixed_string'], np.array([b'fixed_string2'])) self.assertDictElementsAlmostEqual( got[0][constants.FEATURES_KEY][2], { 'fixed_int': np.array([2]), 'fixed_float': np.array([0.0]), }) self.assertEqual( got[0][constants.FEATURES_KEY][2]['fixed_string'], np.array([b'fixed_string3'])) 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]) schema = text_format.Parse( """ tensor_representation_group { key: "" value { tensor_representation { key: "fixed_int" value { dense_tensor { column_name: "fixed_int" } } } } } 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) 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) 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]) self.assertAllClose(got[0][constants.PREDICTIONS_KEY]['model1'], np.array([1.0, 1.0])) self.assertAllClose(got[0][constants.PREDICTIONS_KEY]['model2'], { 'output1': np.array([1.0, 1.0]), 'output2': np.array([0.0, 1.0]) }) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
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) for model in ('model1', 'model2') if multi_model else (''): per_model_result = got[constants.PREDICTIONS_KEY] if model: self.assertIn(model, per_model_result) per_model_result = per_model_result[model] for output in ('Identity', 'Identity_1') if multi_output else (''): per_output_result = per_model_result if output: self.assertIn(output, per_output_result) per_output_result = per_output_result[output] self.assertLen(per_output_result, 2) 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 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 testLabelsExtractor(self, label): model_spec = config.ModelSpec(label_key=label) eval_config = config.EvalConfig(model_specs=[model_spec]) feature_extractor = features_extractor.FeaturesExtractor(eval_config) label_extractor = labels_extractor.LabelsExtractor(eval_config) label_feature = '' if label is not None: label_feature = """ feature { name: "%s" type: FLOAT } """ % label schema = text_format.Parse( label_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, }, label, 1.0), maybe_add_key({ 'fixed_int': 1, }, label, 0.0), maybe_add_key({ 'fixed_int': 2, }, label, 0.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 | label_extractor.stage_name >> label_extractor.ptransform) # pylint: enable=no-value-for-parameter def check_result(got): try: self.assertLen(got, 1) tf.compat.v1.logging.error('HERE >>>> {}'.format(got)) self.assertAlmostEqual( got[0][constants.LABELS_KEY][0], np.array([1.0]) if label is not None else None) self.assertAlmostEqual( got[0][constants.LABELS_KEY][1], np.array([0.0]) if label is not None else None) self.assertAlmostEqual( got[0][constants.LABELS_KEY][2], np.array([0.0]) if label is not None else None) except AssertionError as err: raise util.BeamAssertException(err) util.assert_that(result, check_result, label='result')
def _runMiniPipeline(self, multi_model): """Benchmark a "mini" TFMA - predict, slice and compute metrics. Runs a "mini" version of TFMA in a Beam pipeline. Records the wall time taken for the whole pipeline. Args: multi_model: True if multiple models should be used in the benchmark. """ self._init_model(multi_model, validation=False) pipeline = self._create_beam_pipeline() tfx_io = test_util.InMemoryTFExampleRecord( schema=benchmark_utils.read_schema( self._dataset.tf_metadata_schema_path()), raw_record_column_name=constants.ARROW_INPUT_COLUMN) raw_data = (pipeline | "Examples" >> beam.Create( self._dataset.read_raw_dataset( deserialize=False, limit=self._max_num_examples())) | "BatchExamples" >> tfx_io.BeamSource() | "InputsToExtracts" >> tfma.BatchedInputsToExtracts()) def rescale_labels(extracts): # Transform labels to [0, 1] so we can test metrics that require labels in # that range. result = copy.copy(extracts) result[constants.LABELS_KEY] = self._transform_labels( extracts[constants.LABELS_KEY]) return result _ = (raw_data | "FeaturesExtractor" >> features_extractor.FeaturesExtractor( eval_config=self._eval_config).ptransform | "LabelsExtractor" >> labels_extractor.LabelsExtractor( eval_config=self._eval_config).ptransform | "RescaleLabels" >> beam.Map(rescale_labels) | "ExampleWeightsExtractor" >> example_weights_extractor. ExampleWeightsExtractor(eval_config=self._eval_config).ptransform | "PredictionsExtractor" >> predictions_extractor.PredictionsExtractor( eval_config=self._eval_config, eval_shared_model=self._eval_shared_models).ptransform | "UnbatchExtractor" >> unbatch_extractor.UnbatchExtractor().ptransform | "SliceKeyExtractor" >> tfma.extractors.SliceKeyExtractor().ptransform | "ComputeMetricsPlotsAndValidations" >> metrics_plots_and_validations_evaluator. MetricsPlotsAndValidationsEvaluator( eval_config=self._eval_config, eval_shared_model=self._eval_shared_models).ptransform) start = time.time() for _ in range(_ITERS): result = pipeline.run() result.wait_until_finish() end = time.time() delta = end - start self.report_benchmark( iters=_ITERS, wall_time=delta, extras={ "num_examples": self._dataset.num_examples(limit=self._max_num_examples()) })