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