def testRaggedTensorStructTypeInvalidSteps(self):
tensor_representation = text_format.Parse(
"""
ragged_tensor {
feature_path {
step: "ragged_feature"
step: "wrong_step"
}
}
""", schema_pb2.TensorRepresentation())
record_batch = pa.RecordBatch.from_arrays([
pa.StructArray.from_arrays([
pa.array([[1, 2, 3]], pa.list_(pa.int64())),
pa.array([["a", "b", "c"]], pa.list_(pa.binary()))
], ["inner_feature", "x2"])
], ["ragged_feature"])
with self.assertRaisesRegex(ValueError,
".*Unable to handle tensor output.*"):
tensor_adapter.TensorAdapter(
tensor_adapter.TensorAdapterConfig(
record_batch.schema, {"output": tensor_representation}))
def testRaiseOnRequestingEagerTensorsInGraphMode(self):
tensor_representation = text_format.Parse(
"""
sparse_tensor {
index_column_names: ["key"]
value_column_name: "value"
dense_shape {
dim {
size: 100
}
}
}
""", schema_pb2.TensorRepresentation())
record_batch = pa.RecordBatch.from_arrays(
[pa.array([[1]]), pa.array([[2]])], ["key", "value"])
adapter = tensor_adapter.TensorAdapter(
tensor_adapter.TensorAdapterConfig(
record_batch.schema, {"output": tensor_representation}))
with self.assertRaisesRegex(RuntimeError,
"eager mode was not enabled"):
adapter.ToBatchTensors(record_batch, produce_eager_tensors=True)
def testRaggedTensor(self, tensor_representation_textpb, record_batch,
expected_type_spec, expected_ragged_tensor):
tensor_representation = text_format.Parse(
tensor_representation_textpb, schema_pb2.TensorRepresentation())
adapter = tensor_adapter.TensorAdapter(
tensor_adapter.TensorAdapterConfig(
record_batch.schema, {"output": tensor_representation}))
converted = adapter.ToBatchTensors(record_batch)
self.assertLen(converted, 1)
self.assertIn("output", converted)
actual_output = converted["output"]
self.assertIsInstance(
actual_output,
(tf.RaggedTensor, tf.compat.v1.ragged.RaggedTensorValue))
if tf.executing_eagerly():
self.assertTrue(
expected_type_spec.is_compatible_with(actual_output),
"{} is not compatible with spec {}".format(
actual_output, expected_type_spec))
self.assertRaggedAllEqual(actual_output, expected_ragged_tensor)
self.assertAdapterCanProduceNonEagerInEagerMode(adapter, record_batch)
def testRaggedTensorSlicedRecordBatch(self):
tensor_representation = text_format.Parse(
"""
ragged_tensor {
feature_path {
step: "ragged_feature"
}
}
""", schema_pb2.TensorRepresentation())
record_batch = pa.RecordBatch.from_arrays([
pa.array([[1], None, [2], [3, 4, 5], []],
type=pa.list_(pa.int64()))
], ["ragged_feature"])
record_batch = record_batch.slice(1, 3)
adapter = tensor_adapter.TensorAdapter(
tensor_adapter.TensorAdapterConfig(
record_batch.schema, {"output": tensor_representation}))
with self.assertRaisesRegex(
ValueError,
".*We currently do not handle converting slices to RaggedTensors."
):
adapter.ToBatchTensors(record_batch)
def convert_and_check(tensors, test_values_conversion):
converter = tensor_to_arrow.TensorsToRecordBatchConverter(
type_specs, options)
self.assertEqual(
{f.name: f.type
for f in converter.arrow_schema()}, expected_schema,
"actual: {}".format(converter.arrow_schema()))
canonical_expected_tensor_representations = {}
for n, r in expected_tensor_representations.items():
if not isinstance(r, schema_pb2.TensorRepresentation):
r = text_format.Parse(r, schema_pb2.TensorRepresentation())
canonical_expected_tensor_representations[n] = r
self.assertEqual(canonical_expected_tensor_representations,
converter.tensor_representations())
rb = converter.convert(tensors)
self.assertLen(expected_record_batch, rb.num_columns)
for i, column in enumerate(rb):
expected = expected_record_batch[rb.schema[i].name]
self.assertTrue(
column.equals(expected),
"{}: actual: {}, expected: {}".format(
rb.schema[i].name, column, expected))
# Test that TensorAdapter(TensorsToRecordBatchConverter()) is identity.
adapter = tensor_adapter.TensorAdapter(
tensor_adapter.TensorAdapterConfig(
arrow_schema=converter.arrow_schema(),
tensor_representations=converter.tensor_representations()))
adapter_output = adapter.ToBatchTensors(
rb, produce_eager_tensors=not test_values_conversion)
self.assertEqual(adapter_output.keys(), tensors.keys())
for k in adapter_output.keys():
if "value" not in k:
self._assert_tensor_alike_equal(adapter_output[k],
tensors[k])
def testPassthroughKeys(self):
passthrough_key = '__passthrough__'
def preprocessing_fn(inputs):
self.assertNotIn(passthrough_key, inputs)
return {'x_scaled': tft.scale_to_0_1(inputs['x'])}
x_data = [0., 1., 2.]
passthrough_data = [1, None, 3]
input_record_batch = pa.RecordBatch.from_arrays([
pa.array([[x] for x in x_data], type=pa.list_(pa.float32())),
pa.array([None if p is None else [p] for p in passthrough_data],
type=pa.list_(pa.int64())),
], ['x', passthrough_key])
tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
input_record_batch.schema, {
'x':
text_format.Parse('dense_tensor { column_name: "x" shape {} }',
schema_pb2.TensorRepresentation())
})
expected_data = [{
'x_scaled': x / 2.0,
passthrough_key: p
} for x, p in zip(x_data, passthrough_data)]
with self._makeTestPipeline() as pipeline:
input_data = (pipeline | beam.Create([input_record_batch]))
with beam_impl.Context(temp_dir=self.get_temp_dir(),
passthrough_keys=set([passthrough_key])):
(transformed_data, _), _ = (
(input_data, tensor_adapter_config)
| beam_impl.AnalyzeAndTransformDataset(preprocessing_fn))
def _assert_fn(output_data):
self.assertCountEqual(expected_data, output_data)
beam_test_util.assert_that(transformed_data, _assert_fn)
def setUp(self):
super().setUp()
self._eval_export_dir = os.path.join(self._getTempDir(), 'eval_export')
self._create_sklearn_model(self._eval_export_dir)
self._eval_config = tfma.EvalConfig(model_specs=[tfma.ModelSpec()])
self._eval_shared_model = (
sklearn_predict_extractor.custom_eval_shared_model(
eval_saved_model_path=self._eval_export_dir,
model_name=None,
eval_config=self._eval_config))
self._schema = text_format.Parse(
"""
feature {
name: "age"
type: FLOAT
}
feature {
name: "language"
type: FLOAT
}
feature {
name: "label"
type: INT
}
""", schema_pb2.Schema())
self._tfx_io = test_util.InMemoryTFExampleRecord(
schema=self._schema,
raw_record_column_name=tfma.ARROW_INPUT_COLUMN)
self._tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
arrow_schema=self._tfx_io.ArrowSchema(),
tensor_representations=self._tfx_io.TensorRepresentations())
self._examples = [
self._makeExample(age=3.0, language=1.0, label=1),
self._makeExample(age=3.0, language=0.0, label=0),
self._makeExample(age=4.0, language=1.0, label=1),
self._makeExample(age=5.0, language=0.0, label=0),
]
def test_convert(self, type_specs, expected_schema,
expected_tensor_representations, tensor_input,
expected_record_batch):
converter = tensor_to_arrow.TensorsToRecordBatchConverter(type_specs)
expected_schema = pa.schema(
[pa.field(n, t) for n, t in sorted(expected_schema.items())])
self.assertTrue(converter.arrow_schema().equals(expected_schema),
"actual: {}".format(converter.arrow_schema()))
canonical_expected_tensor_representations = {}
for n, r in expected_tensor_representations.items():
if not isinstance(r, schema_pb2.TensorRepresentation):
r = text_format.Parse(r, schema_pb2.TensorRepresentation())
canonical_expected_tensor_representations[n] = r
self.assertEqual(canonical_expected_tensor_representations,
converter.tensor_representations())
rb = converter.convert(tensor_input)
self.assertTrue(
rb.equals(
pa.record_batch(
[arr for _, arr in sorted(expected_record_batch.items())],
schema=expected_schema)))
# Test that TensorAdapter(TensorsToRecordBatchConverter()) is identity.
adapter = tensor_adapter.TensorAdapter(
tensor_adapter.TensorAdapterConfig(
arrow_schema=converter.arrow_schema(),
tensor_representations=converter.tensor_representations()))
adapter_output = adapter.ToBatchTensors(rb, produce_eager_tensors=True)
self.assertEqual(adapter_output.keys(), tensor_input.keys())
for k in adapter_output.keys():
self._assert_tensor_alike_equal(adapter_output[k], tensor_input[k])
def testOneTensorFromOneColumn(self, tensor_representation_textpb,
arrow_array, expected_type_spec,
expected_output):
tensor_representation = text_format.Parse(
tensor_representation_textpb, schema_pb2.TensorRepresentation())
column_name = None
if tensor_representation.HasField("dense_tensor"):
column_name = tensor_representation.dense_tensor.column_name
if tensor_representation.HasField("varlen_sparse_tensor"):
column_name = tensor_representation.varlen_sparse_tensor.column_name
record_batch = pa.RecordBatch.from_arrays([arrow_array], [column_name])
adapter = tensor_adapter.TensorAdapter(
tensor_adapter.TensorAdapterConfig(
record_batch.schema, {"output": tensor_representation}))
self.assertEqual(expected_type_spec, adapter.TypeSpecs()["output"])
converted = adapter.ToBatchTensors(record_batch)
self.assertLen(converted, 1)
self.assertIn("output", converted)
actual_output = converted["output"]
if tf.executing_eagerly():
self.assertTrue(
expected_type_spec.is_compatible_with(actual_output),
"{} is not compatible with spec {}".format(
actual_output, expected_type_spec))
if isinstance(expected_output,
(tf.SparseTensor, tf.compat.v1.SparseTensorValue)):
self.assertIsInstance(
actual_output,
(tf.SparseTensor, tf.compat.v1.SparseTensorValue))
self.assertSparseAllEqual(expected_output, actual_output)
else:
self.assertAllEqual(expected_output, actual_output)
self.assertAdapterCanProduceNonEagerInEagerMode(adapter, record_batch)
def record_batch_to_tensor_values(
record_batch: pa.RecordBatch,
tensor_representations: Optional[Mapping[
str, schema_pb2.TensorRepresentation]] = None
) -> types.TensorValueMaybeMultiLevelDict:
"""Returns tensor values extracted from given record batch.
Args:
record_batch: Record batch to extract features from.
tensor_representations: Tensor representations to use when extracting the
features. If a representation is not found for a given column name, a
default representation will be used where possible, otherwise an exception
will be raised.
Returns:
Features dict.
Raises:
ValueError: If a tensor value cannot be determined for a given column in the
record batch.
"""
if tensor_representations is None:
tensor_representations = {}
def _shape(value: Any) -> List[int]:
"""Returns the shape associated with given value."""
if hasattr(value, '__len__'):
return [len(value)] + _shape(value[0]) if value else [len(value)]
else:
return []
features = {}
updated_tensor_representations = {}
for i, col in enumerate(record_batch.schema):
if col.name in tensor_representations:
updated_tensor_representations[col.name] = (
tensor_representations[col.name])
else:
col_sizes = record_batch.column(i).value_lengths().unique()
if len(col_sizes) != 1:
# Assume VarLenSparseTensor
tensor_representation = schema_pb2.TensorRepresentation()
tensor_representation.varlen_sparse_tensor.column_name = col.name
updated_tensor_representations[
col.name] = tensor_representation
elif not np.all(record_batch[i].is_valid()):
# Features that are missing some values can't be parsed using a default
# tensor representation. Convert to numpy arrays containing None values.
features[col.name] = record_batch[i].to_numpy(
zero_copy_only=False)
else:
tensor_representation = schema_pb2.TensorRepresentation()
tensor_representation.dense_tensor.column_name = col.name
dims = _shape(record_batch[i])
# Convert dims of the form (..., n, 1) to (..., n).
if len(dims) > 1 and dims[-1] == 1:
dims = dims[:-1]
if len(dims) > 1:
for dim in dims[1:]: # Skip batch dimension
tensor_representation.dense_tensor.shape.dim.append(
schema_pb2.FixedShape.Dim(size=dim))
updated_tensor_representations[
col.name] = tensor_representation
if updated_tensor_representations:
adapter = tensor_adapter.TensorAdapter(
tensor_adapter.TensorAdapterConfig(
arrow_schema=record_batch.schema,
tensor_representations=updated_tensor_representations))
try:
for k, v in adapter.ToBatchTensors(
record_batch, produce_eager_tensors=False).items():
if isinstance(v, tf.compat.v1.ragged.RaggedTensorValue):
features[k] = to_ragged_tensor_value(v)
elif isinstance(v, tf.compat.v1.SparseTensorValue):
kind = updated_tensor_representations[k].WhichOneof('kind')
if kind == 'sparse_tensor':
features[k] = to_sparse_tensor_value(v)
elif kind == 'varlen_sparse_tensor':
features[k] = to_varlen_sparse_tensor_value(v)
else:
raise ValueError(
f'Unexpected tensor representation kind ({kind}) '
f'for tensor of type: {type(v)}')
else:
features[k] = v
except Exception as e:
raise ValueError(e, updated_tensor_representations,
record_batch) from e
return features
def Do(self, input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]) -> None:
# Check the inputs
if constants.EXAMPLES not in input_dict:
raise ValueError(f'{constants.EXAMPLES} is missing from inputs')
examples_artifact = input_dict[constants.EXAMPLES]
input_uri = artifact_utils.get_single_uri(examples_artifact)
if len(zenml_path_utils.list_dir(input_uri)) == 0:
raise AssertionError(
'ZenML can not run the evaluation as the provided input '
'configuration does not point towards any data. Specifically, '
'if you are using the agnostic evaluator, please make sure '
'that you are using a proper test_fn in your trainer step to '
'write these results.')
else:
# Check the outputs
if constants.EVALUATION not in output_dict:
raise ValueError(
f'{constants.EVALUATION} is missing from outputs')
evaluation_artifact = output_dict[constants.EVALUATION]
output_uri = artifact_utils.get_single_uri(evaluation_artifact)
# Resolve the schema
schema = None
if constants.SCHEMA in input_dict:
schema_artifact = input_dict[constants.SCHEMA]
schema_uri = artifact_utils.get_single_uri(schema_artifact)
reader = io_utils.SchemaReader()
schema = reader.read(io_utils.get_only_uri_in_dir(schema_uri))
# Create the step with the schema attached if provided
source = exec_properties[StepKeys.SOURCE]
args = exec_properties[StepKeys.ARGS]
c = source_utils.load_source_path_class(source)
evaluator_step: BaseEvaluatorStep = c(**args)
# Check the execution parameters
eval_config = evaluator_step.build_config()
eval_config = tfma.update_eval_config_with_defaults(eval_config)
tfma.verify_eval_config(eval_config)
# Resolve the model
if constants.MODEL in input_dict:
model_artifact = input_dict[constants.MODEL]
model_uri = artifact_utils.get_single_uri(model_artifact)
model_path = path_utils.serving_model_path(model_uri)
model_fn = try_get_fn(evaluator_step.CUSTOM_MODULE,
'custom_eval_shared_model'
) or tfma.default_eval_shared_model
eval_shared_model = model_fn(
model_name='', # TODO: Fix with model names
eval_saved_model_path=model_path,
eval_config=eval_config)
else:
eval_shared_model = None
self._log_startup(input_dict, output_dict, exec_properties)
# Main pipeline
logging.info('Evaluating model.')
with self._make_beam_pipeline() as pipeline:
examples_list = []
tensor_adapter_config = None
if tfma.is_batched_input(eval_shared_model, eval_config):
tfxio_factory = tfxio_utils.get_tfxio_factory_from_artifact(
examples=[
artifact_utils.get_single_instance(
examples_artifact)
],
telemetry_descriptors=_TELEMETRY_DESCRIPTORS,
schema=schema,
raw_record_column_name=tfma_constants.
ARROW_INPUT_COLUMN)
for split in evaluator_step.splits:
file_pattern = io_utils.all_files_pattern(
artifact_utils.get_split_uri(
examples_artifact, split))
tfxio = tfxio_factory(file_pattern)
data = (pipeline
| 'ReadFromTFRecordToArrow[%s]' % split >>
tfxio.BeamSource())
examples_list.append(data)
if schema is not None:
tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
arrow_schema=tfxio.ArrowSchema(),
tensor_representations=tfxio.TensorRepresentations(
))
else:
for split in evaluator_step.splits:
file_pattern = io_utils.all_files_pattern(
artifact_utils.get_split_uri(
examples_artifact, split))
data = (pipeline
| 'ReadFromTFRecord[%s]' % split >> beam.io.
ReadFromTFRecord(file_pattern=file_pattern))
examples_list.append(data)
# Resolve custom extractors
custom_extractors = try_get_fn(evaluator_step.CUSTOM_MODULE,
'custom_extractors')
extractors = None
if custom_extractors:
extractors = custom_extractors(
eval_shared_model=eval_shared_model,
eval_config=eval_config,
tensor_adapter_config=tensor_adapter_config)
# Resolve custom evaluators
custom_evaluators = try_get_fn(evaluator_step.CUSTOM_MODULE,
'custom_evaluators')
evaluators = None
if custom_evaluators:
evaluators = custom_evaluators(
eval_shared_model=eval_shared_model,
eval_config=eval_config,
tensor_adapter_config=tensor_adapter_config)
# Extract, evaluate and write
(examples_list | 'FlattenExamples' >> beam.Flatten()
| 'ExtractEvaluateAndWriteResults' >>
tfma.ExtractEvaluateAndWriteResults(
eval_config=eval_config,
eval_shared_model=eval_shared_model,
output_path=output_uri,
extractors=extractors,
evaluators=evaluators,
tensor_adapter_config=tensor_adapter_config))
logging.info('Evaluation complete. Results written to %s.',
output_uri)
def append_tfma_pipeline(pipeline: beam.Pipeline,
me_eval_config: me_proto.EvaluationConfig,
problem_type: constants.ProblemType,
tfma_format: Optional[bool] = False,
json_mode: Optional[bool] = False,
schema: Optional[Any] = None):
"""Extend a beam pipeline to add TFMA evaluation given a configuration.
Args:
pipeline: A beam pipeline.
me_eval_config: A ME Evaluation Configuration.
problem_type: Defines what type of problem to expect.
tfma_format: If true, use TFMA format, if false use Model Evaluation.
json_mode: Output metrics in a plain text mode.
schema: Optional tf.metadata schema. If you need to pass multi-tensor input
to the model, you need to pass the schema.
"""
input_files = (
me_eval_config.data_spec.input_source_spec.jsonl_file_spec.file_names)
output_path = me_eval_config.output_spec.gcs_sink.path
data_spec = me_eval_config.data_spec
weight_column_spec = ColumnSpec(
me_eval_config.data_spec.example_weight_key_spec
) if me_eval_config.data_spec.HasField('example_weight_key_spec') else None
eval_column_specs = EvaluationColumnSpecs(
ground_truth_column_spec=ColumnSpec(
me_eval_config.data_spec.label_key_spec),
example_weight_column_spec=weight_column_spec,
predicted_score_column_spec=ColumnSpec(
data_spec.predicted_score_key_spec)
if data_spec.HasField('predicted_score_key_spec') else None,
predicted_label_column_spec=ColumnSpec(
data_spec.predicted_label_key_spec)
if data_spec.HasField('predicted_label_key_spec') else None,
predicted_label_id_column_spec=ColumnSpec(
data_spec.predicted_label_id_key_spec)
if data_spec.HasField('predicted_label_id_key_spec') else None)
class_name_list = list(me_eval_config.data_spec.labels) or None
quantile_list = list(data_spec.quantiles) or None
quantile_index = data_spec.quantile_index if data_spec.quantile_index >= 0 else None
tfma_eval_config = tfma_adapter.METoTFMA(class_name_list).eval_config(
me_eval_config)
me_writers = [
tfma.writers.Writer(
stage_name='WriteMetrics',
# pylint:disable=no-value-for-parameter
ptransform=_write_metrics(output_file=os.path.join(
output_path, constants.Pipeline.METRICS_KEY),
problem_type=problem_type,
class_labels=class_name_list,
tfma_format=tfma_format,
json_mode=json_mode)),
]
coder = tf_example_record.TFExampleBeamRecord(
physical_format='inmem',
schema=schema,
raw_record_column_name=tfma.ARROW_INPUT_COLUMN,
telemetry_descriptors=None)
tensor_adapter_config = None
if schema is not None:
tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
arrow_schema=coder.ArrowSchema(),
tensor_representations=coder.TensorRepresentations())
_ = (pipeline | 'InputFileList' >> beam.Create(input_files)
| 'ReadText' >> beam.io.textio.ReadAllFromText()
| 'ParseData' >> beam.ParDo(
JSONToSerializedExample(eval_column_specs=eval_column_specs,
class_list=class_name_list,
quantile_list=quantile_list,
quantile_index=quantile_index))
| 'ExamplesToRecordBatch' >> coder.BeamSource()
| 'ExtractEvaluateAndWriteResults' >>
tfma.ExtractEvaluateAndWriteResults(
eval_config=tfma_eval_config,
writers=me_writers,
tensor_adapter_config=tensor_adapter_config))
def testRaiseOnUnsupportedTensorRepresentation(self):
with self.assertRaisesRegex(ValueError, "Unable to handle tensor"):
tensor_adapter.TensorAdapter(
tensor_adapter.TensorAdapterConfig(
pa.schema([pa.field("a", pa.list_(pa.int64()))]),
{"tensor": schema_pb2.TensorRepresentation()}))
def Do(self, input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]) -> None:
"""Runs a batch job to evaluate the eval_model against the given input.
Args:
input_dict: Input dict from input key to a list of Artifacts.
- model_exports: exported model.
- examples: examples for eval the model.
output_dict: Output dict from output key to a list of Artifacts.
- output: model evaluation results.
exec_properties: A dict of execution properties.
- eval_config: JSON string of tfma.EvalConfig.
- feature_slicing_spec: JSON string of evaluator_pb2.FeatureSlicingSpec
instance, providing the way to slice the data. Deprecated, use
eval_config.slicing_specs instead.
- example_splits: JSON-serialized list of names of splits on which the
metrics are computed. Default behavior (when example_splits is set to
None) is using the 'eval' split.
Returns:
None
"""
if constants.EXAMPLES_KEY not in input_dict:
raise ValueError('EXAMPLES_KEY is missing from input dict.')
if constants.MODEL_KEY not in input_dict:
raise ValueError('MODEL_KEY is missing from input dict.')
if constants.EVALUATION_KEY not in output_dict:
raise ValueError('EVALUATION_KEY is missing from output dict.')
if len(input_dict[constants.MODEL_KEY]) > 1:
raise ValueError(
'There can be only one candidate model, there are %d.' %
(len(input_dict[constants.MODEL_KEY])))
if constants.BASELINE_MODEL_KEY in input_dict and len(
input_dict[constants.BASELINE_MODEL_KEY]) > 1:
raise ValueError(
'There can be only one baseline model, there are %d.' %
(len(input_dict[constants.BASELINE_MODEL_KEY])))
self._log_startup(input_dict, output_dict, exec_properties)
# Add fairness indicator metric callback if necessary.
fairness_indicator_thresholds = exec_properties.get(
'fairness_indicator_thresholds', None)
add_metrics_callbacks = None
if fairness_indicator_thresholds:
add_metrics_callbacks = [
tfma.post_export_metrics.fairness_indicators( # pytype: disable=module-attr
thresholds=fairness_indicator_thresholds),
]
output_uri = artifact_utils.get_single_uri(
output_dict[constants.EVALUATION_KEY])
eval_shared_model_fn = udf_utils.try_get_fn(
exec_properties=exec_properties,
fn_name='custom_eval_shared_model'
) or tfma.default_eval_shared_model
run_validation = False
models = []
if 'eval_config' in exec_properties and exec_properties['eval_config']:
slice_spec = None
has_baseline = bool(input_dict.get(constants.BASELINE_MODEL_KEY))
eval_config = tfma.EvalConfig()
json_format.Parse(exec_properties['eval_config'], eval_config)
eval_config = tfma.update_eval_config_with_defaults(
eval_config,
maybe_add_baseline=has_baseline,
maybe_remove_baseline=not has_baseline)
tfma.verify_eval_config(eval_config)
# Do not validate model when there is no thresholds configured. This is to
# avoid accidentally blessing models when users forget to set thresholds.
run_validation = bool(
tfma.metrics.metric_thresholds_from_metrics_specs(
eval_config.metrics_specs))
if len(eval_config.model_specs) > 2:
raise ValueError(
"""Cannot support more than two models. There are %d models in this
eval_config.""" % (len(eval_config.model_specs)))
# Extract model artifacts.
for model_spec in eval_config.model_specs:
if model_spec.is_baseline:
model_uri = artifact_utils.get_single_uri(
input_dict[constants.BASELINE_MODEL_KEY])
else:
model_uri = artifact_utils.get_single_uri(
input_dict[constants.MODEL_KEY])
if tfma.get_model_type(model_spec) == tfma.TF_ESTIMATOR:
model_path = path_utils.eval_model_path(model_uri)
else:
model_path = path_utils.serving_model_path(model_uri)
logging.info('Using %s as %s model.', model_path,
model_spec.name)
models.append(
eval_shared_model_fn(
eval_saved_model_path=model_path,
model_name=model_spec.name,
eval_config=eval_config,
add_metrics_callbacks=add_metrics_callbacks))
else:
eval_config = None
assert ('feature_slicing_spec' in exec_properties
and exec_properties['feature_slicing_spec']
), 'both eval_config and feature_slicing_spec are unset.'
feature_slicing_spec = evaluator_pb2.FeatureSlicingSpec()
json_format.Parse(exec_properties['feature_slicing_spec'],
feature_slicing_spec)
slice_spec = self._get_slice_spec_from_feature_slicing_spec(
feature_slicing_spec)
model_uri = artifact_utils.get_single_uri(
input_dict[constants.MODEL_KEY])
model_path = path_utils.eval_model_path(model_uri)
logging.info('Using %s for model eval.', model_path)
models.append(
eval_shared_model_fn(
eval_saved_model_path=model_path,
model_name='',
eval_config=None,
add_metrics_callbacks=add_metrics_callbacks))
eval_shared_model = models[0] if len(models) == 1 else models
schema = None
if constants.SCHEMA_KEY in input_dict:
schema = io_utils.SchemaReader().read(
io_utils.get_only_uri_in_dir(
artifact_utils.get_single_uri(
input_dict[constants.SCHEMA_KEY])))
# Load and deserialize example splits from execution properties.
example_splits = json_utils.loads(
exec_properties.get(constants.EXAMPLE_SPLITS_KEY, 'null'))
if not example_splits:
example_splits = ['eval']
logging.info(
"The 'example_splits' parameter is not set, using 'eval' "
'split.')
logging.info('Evaluating model.')
with self._make_beam_pipeline() as pipeline:
examples_list = []
tensor_adapter_config = None
# pylint: disable=expression-not-assigned
if _USE_TFXIO and tfma.is_batched_input(eval_shared_model,
eval_config):
tfxio_factory = tfxio_utils.get_tfxio_factory_from_artifact(
examples=[
artifact_utils.get_single_instance(
input_dict[constants.EXAMPLES_KEY])
],
telemetry_descriptors=_TELEMETRY_DESCRIPTORS,
schema=schema,
raw_record_column_name=tfma_constants.ARROW_INPUT_COLUMN)
# TODO(b/161935932): refactor after TFXIO supports multiple patterns.
for split in example_splits:
file_pattern = io_utils.all_files_pattern(
artifact_utils.get_split_uri(
input_dict[constants.EXAMPLES_KEY], split))
tfxio = tfxio_factory(file_pattern)
data = (pipeline
| 'ReadFromTFRecordToArrow[%s]' % split >>
tfxio.BeamSource())
examples_list.append(data)
if schema is not None:
# Use last tfxio as TensorRepresentations and ArrowSchema are fixed.
tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
arrow_schema=tfxio.ArrowSchema(),
tensor_representations=tfxio.TensorRepresentations())
else:
for split in example_splits:
file_pattern = io_utils.all_files_pattern(
artifact_utils.get_split_uri(
input_dict[constants.EXAMPLES_KEY], split))
data = (
pipeline
| 'ReadFromTFRecord[%s]' % split >>
beam.io.ReadFromTFRecord(file_pattern=file_pattern))
examples_list.append(data)
custom_extractors = udf_utils.try_get_fn(
exec_properties=exec_properties, fn_name='custom_extractors')
extractors = None
if custom_extractors:
extractors = custom_extractors(
eval_shared_model=eval_shared_model,
eval_config=eval_config,
tensor_adapter_config=tensor_adapter_config)
(examples_list | 'FlattenExamples' >> beam.Flatten()
| 'ExtractEvaluateAndWriteResults' >>
tfma.ExtractEvaluateAndWriteResults(
eval_shared_model=models[0] if len(models) == 1 else models,
eval_config=eval_config,
extractors=extractors,
output_path=output_uri,
slice_spec=slice_spec,
tensor_adapter_config=tensor_adapter_config))
logging.info('Evaluation complete. Results written to %s.', output_uri)
if not run_validation:
# TODO(jinhuang): delete the BLESSING_KEY from output_dict when supported.
logging.info('No threshold configured, will not validate model.')
return
# Set up blessing artifact
blessing = artifact_utils.get_single_instance(
output_dict[constants.BLESSING_KEY])
blessing.set_string_custom_property(
constants.ARTIFACT_PROPERTY_CURRENT_MODEL_URI_KEY,
artifact_utils.get_single_uri(input_dict[constants.MODEL_KEY]))
blessing.set_int_custom_property(
constants.ARTIFACT_PROPERTY_CURRENT_MODEL_ID_KEY,
input_dict[constants.MODEL_KEY][0].id)
if input_dict.get(constants.BASELINE_MODEL_KEY):
baseline_model = input_dict[constants.BASELINE_MODEL_KEY][0]
blessing.set_string_custom_property(
constants.ARTIFACT_PROPERTY_BASELINE_MODEL_URI_KEY,
baseline_model.uri)
blessing.set_int_custom_property(
constants.ARTIFACT_PROPERTY_BASELINE_MODEL_ID_KEY,
baseline_model.id)
if 'current_component_id' in exec_properties:
blessing.set_string_custom_property(
'component_id', exec_properties['current_component_id'])
# Check validation result and write BLESSED file accordingly.
logging.info('Checking validation results.')
validation_result = tfma.load_validation_result(output_uri)
if validation_result.validation_ok:
io_utils.write_string_file(
os.path.join(blessing.uri, constants.BLESSED_FILE_NAME), '')
blessing.set_int_custom_property(
constants.ARTIFACT_PROPERTY_BLESSED_KEY,
constants.BLESSED_VALUE)
else:
io_utils.write_string_file(
os.path.join(blessing.uri, constants.NOT_BLESSED_FILE_NAME),
'')
blessing.set_int_custom_property(
constants.ARTIFACT_PROPERTY_BLESSED_KEY,
constants.NOT_BLESSED_VALUE)
logging.info('Blessing result %s written to %s.',
validation_result.validation_ok, blessing.uri)
def TensorAdapterConfig(self) -> tensor_adapter.TensorAdapterConfig:
return tensor_adapter.TensorAdapterConfig(
self.projected.ArrowSchema(),
self.projected.TensorRepresentations(),
original_type_specs=self.origin.TensorAdapter().TypeSpecs())
def testMultipleColumns(self):
record_batch = pa.RecordBatch.from_arrays([
pa.array([[1], [], [2, 3], None], type=pa.large_list(pa.int64())),
pa.array([[1.0, 2.0], [2.0, 3.0], [3.0, 4.0], [4.0, 5.0]],
type=pa.list_(pa.float32())),
pa.array([None, [b"a", b"b"], [b"c", b"d"], None],
type=pa.list_(pa.large_binary())),
pa.array([[b"w"], [b"x"], [b"y"], [b"z"]],
type=pa.list_(pa.string())),
], [
"int64_ragged",
"float_dense",
"bytes_ragged",
"bytes_dense",
])
tensor_representations = {
"int64_varlen_sparse":
text_format.Parse(
"""
varlen_sparse_tensor {
column_name: "int64_ragged"
}
""", schema_pb2.TensorRepresentation()),
"float_dense":
text_format.Parse(
"""
dense_tensor {
column_name: "float_dense"
shape {
dim {
size: 2
}
dim {
size: 1
}
}
}""", schema_pb2.TensorRepresentation()),
"bytes_varlen_sparse":
text_format.Parse(
"""
varlen_sparse_tensor {
column_name: "bytes_ragged"
}
""", schema_pb2.TensorRepresentation()),
"bytes_dense":
text_format.Parse(
"""
dense_tensor {
column_name: "bytes_dense"
shape {
}
}
""", schema_pb2.TensorRepresentation()),
"bytes_default_filled_dense":
text_format.Parse(
"""
dense_tensor {
column_name: "bytes_ragged"
shape {
dim {
size: 2
}
}
default_value {
bytes_value: "kk"
}
}
""", schema_pb2.TensorRepresentation()),
}
adapter = tensor_adapter.TensorAdapter(
tensor_adapter.TensorAdapterConfig(record_batch.schema,
tensor_representations))
type_specs = adapter.TypeSpecs()
self.assertEqual(
type_specs, {
"int64_varlen_sparse":
tf.SparseTensorSpec(shape=[None, None], dtype=tf.int64),
"bytes_varlen_sparse":
tf.SparseTensorSpec(shape=[None, None], dtype=tf.string),
"float_dense":
tf.TensorSpec(shape=[None, 2, 1], dtype=tf.float32),
"bytes_dense":
tf.TensorSpec(shape=[None], dtype=tf.string),
"bytes_default_filled_dense":
tf.TensorSpec(shape=[None, 2], dtype=tf.string),
})
tensors = adapter.ToBatchTensors(record_batch)
self.assertLen(tensors, len(type_specs))
self.assertSparseAllEqual(
tf.SparseTensor(values=tf.constant([1, 2, 3], dtype=tf.int64),
dense_shape=tf.constant([4, 2], dtype=tf.int64),
indices=tf.constant([[0, 0], [2, 0], [2, 1]],
dtype=tf.int64)),
tensors["int64_varlen_sparse"])
self.assertSparseAllEqual(
tf.SparseTensor(values=tf.constant([b"a", b"b", b"c", b"d"]),
dense_shape=tf.constant([4, 2], dtype=tf.int64),
indices=tf.constant(
[[1, 0], [1, 1], [2, 0], [2, 1]],
dtype=tf.int64)),
tensors["bytes_varlen_sparse"])
self.assertAllEqual(
tf.constant([[[1.0], [2.0]], [[2.0], [3.0]], [[3.0], [4.0]],
[[4.0], [5.0]]],
dtype=tf.float32), tensors["float_dense"])
self.assertAllEqual(tf.constant([b"w", b"x", b"y", b"z"]),
tensors["bytes_dense"])
self.assertAllEqual(
tf.constant([[b"kk", b"kk"], [b"a", b"b"], [b"c", b"d"],
[b"kk", b"kk"]]),
tensors["bytes_default_filled_dense"])
if tf.executing_eagerly():
for name, spec in six.iteritems(type_specs):
self.assertTrue(
spec.is_compatible_with(tensors[name]),
"{} is not compatible with spec {}".format(
tensors[name], spec))
self.assertAdapterCanProduceNonEagerInEagerMode(adapter, record_batch)
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.BATCHED_INPUT_KEY)
tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
arrow_schema=tfx_io.ArrowSchema(),
tensor_representations=tfx_io.TensorRepresentations())
input_extractor = batched_input_extractor.BatchedInputExtractor(
eval_config)
predict_extractor = batched_predict_extractor_v2.BatchedPredictExtractor(
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()
| input_extractor.stage_name >> input_extractor.ptransform
| predict_extractor.stage_name >> predict_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.BATCHED_PREDICTIONS_KEY, item)
except AssertionError as err:
raise util.BeamAssertException(err)
util.assert_that(predict_extracts, check_result, label='result')
def Do(self, input_dict: Dict[str, List[types.Artifact]],
output_dict: Dict[str, List[types.Artifact]],
exec_properties: Dict[str, Any]) -> None:
"""Runs a batch job to evaluate the eval_model against the given input.
Args:
input_dict: Input dict from input key to a list of Artifacts.
- model: exported model.
- examples: examples for eval the model.
output_dict: Output dict from output key to a list of Artifacts.
- evaluation: model evaluation results.
exec_properties: A dict of execution properties.
- eval_config: JSON string of tfma.EvalConfig.
- feature_slicing_spec: JSON string of evaluator_pb2.FeatureSlicingSpec
instance, providing the way to slice the data. Deprecated, use
eval_config.slicing_specs instead.
- example_splits: JSON-serialized list of names of splits on which the
metrics are computed. Default behavior (when example_splits is set to
None) is using the 'eval' split.
Returns:
None
"""
if standard_component_specs.EXAMPLES_KEY not in input_dict:
raise ValueError('EXAMPLES_KEY is missing from input dict.')
if standard_component_specs.EVALUATION_KEY not in output_dict:
raise ValueError('EVALUATION_KEY is missing from output dict.')
if standard_component_specs.MODEL_KEY in input_dict and len(
input_dict[standard_component_specs.MODEL_KEY]) > 1:
raise ValueError('There can be only one candidate model, there are %d.' %
(len(input_dict[standard_component_specs.MODEL_KEY])))
if standard_component_specs.BASELINE_MODEL_KEY in input_dict and len(
input_dict[standard_component_specs.BASELINE_MODEL_KEY]) > 1:
raise ValueError(
'There can be only one baseline model, there are %d.' %
(len(input_dict[standard_component_specs.BASELINE_MODEL_KEY])))
self._log_startup(input_dict, output_dict, exec_properties)
# Add fairness indicator metric callback if necessary.
fairness_indicator_thresholds = json_utils.loads(
exec_properties.get(
standard_component_specs.FAIRNESS_INDICATOR_THRESHOLDS_KEY, 'null'))
add_metrics_callbacks = None
if fairness_indicator_thresholds:
add_metrics_callbacks = [
tfma.post_export_metrics.fairness_indicators( # pytype: disable=module-attr
thresholds=fairness_indicator_thresholds),
]
output_uri = artifact_utils.get_single_uri(
output_dict[constants.EVALUATION_KEY])
# Make sure user packages get propagated to the remote Beam worker.
unused_module_path, extra_pip_packages = udf_utils.decode_user_module_key(
exec_properties.get(standard_component_specs.MODULE_PATH_KEY, None))
for pip_package_path in extra_pip_packages:
local_pip_package_path = io_utils.ensure_local(pip_package_path)
self._beam_pipeline_args.append('--extra_package=%s' %
local_pip_package_path)
eval_shared_model_fn = udf_utils.try_get_fn(
exec_properties=exec_properties,
fn_name='custom_eval_shared_model') or tfma.default_eval_shared_model
run_validation = False
models = []
if (standard_component_specs.EVAL_CONFIG_KEY in exec_properties
and exec_properties[standard_component_specs.EVAL_CONFIG_KEY]):
slice_spec = None
has_baseline = bool(
input_dict.get(standard_component_specs.BASELINE_MODEL_KEY))
eval_config = tfma.EvalConfig()
proto_utils.json_to_proto(
exec_properties[standard_component_specs.EVAL_CONFIG_KEY],
eval_config)
# rubber_stamp is always assumed true, i.e., change threshold will always
# be ignored when a baseline model is missing.
if hasattr(tfma, 'utils'):
eval_config = tfma.utils.update_eval_config_with_defaults(
eval_config, has_baseline=has_baseline, rubber_stamp=True)
tfma.utils.verify_eval_config(eval_config)
else:
# TODO(b/171992041): Replaced by tfma.utils.
eval_config = tfma.update_eval_config_with_defaults(
eval_config, has_baseline=has_baseline, rubber_stamp=True)
tfma.verify_eval_config(eval_config)
# Do not validate model when there is no thresholds configured. This is to
# avoid accidentally blessing models when users forget to set thresholds.
run_validation = bool(
tfma.metrics.metric_thresholds_from_metrics_specs(
eval_config.metrics_specs, eval_config=eval_config))
if len(eval_config.model_specs) > 2:
raise ValueError(
"""Cannot support more than two models. There are %d models in this
eval_config.""" % (len(eval_config.model_specs)))
# Extract model artifacts.
for model_spec in eval_config.model_specs:
if standard_component_specs.MODEL_KEY not in input_dict:
if not model_spec.prediction_key:
raise ValueError(
'model_spec.prediction_key required if model not provided')
continue
if model_spec.is_baseline:
model_artifact = artifact_utils.get_single_instance(
input_dict[standard_component_specs.BASELINE_MODEL_KEY])
else:
model_artifact = artifact_utils.get_single_instance(
input_dict[standard_component_specs.MODEL_KEY])
# TODO(b/171992041): tfma.get_model_type replaced by tfma.utils.
if ((hasattr(tfma, 'utils') and
tfma.utils.get_model_type(model_spec) == tfma.TF_ESTIMATOR) or
hasattr(tfma, 'get_model_type') and
tfma.get_model_type(model_spec) == tfma.TF_ESTIMATOR):
model_path = path_utils.eval_model_path(
model_artifact.uri,
path_utils.is_old_model_artifact(model_artifact))
else:
model_path = path_utils.serving_model_path(
model_artifact.uri,
path_utils.is_old_model_artifact(model_artifact))
logging.info('Using %s as %s model.', model_path, model_spec.name)
models.append(
eval_shared_model_fn(
eval_saved_model_path=model_path,
model_name=model_spec.name,
eval_config=eval_config,
add_metrics_callbacks=add_metrics_callbacks))
else:
eval_config = None
assert (standard_component_specs.FEATURE_SLICING_SPEC_KEY
in exec_properties and
exec_properties[standard_component_specs.FEATURE_SLICING_SPEC_KEY]
), 'both eval_config and feature_slicing_spec are unset.'
feature_slicing_spec = evaluator_pb2.FeatureSlicingSpec()
proto_utils.json_to_proto(
exec_properties[standard_component_specs.FEATURE_SLICING_SPEC_KEY],
feature_slicing_spec)
slice_spec = self._get_slice_spec_from_feature_slicing_spec(
feature_slicing_spec)
model_artifact = artifact_utils.get_single_instance(
input_dict[standard_component_specs.MODEL_KEY])
model_path = path_utils.eval_model_path(
model_artifact.uri, path_utils.is_old_model_artifact(model_artifact))
logging.info('Using %s for model eval.', model_path)
models.append(
eval_shared_model_fn(
eval_saved_model_path=model_path,
model_name='',
eval_config=None,
add_metrics_callbacks=add_metrics_callbacks))
eval_shared_model = models[0] if len(models) == 1 else models
schema = None
if standard_component_specs.SCHEMA_KEY in input_dict:
schema = io_utils.SchemaReader().read(
io_utils.get_only_uri_in_dir(
artifact_utils.get_single_uri(
input_dict[standard_component_specs.SCHEMA_KEY])))
# Load and deserialize example splits from execution properties.
example_splits = json_utils.loads(
exec_properties.get(standard_component_specs.EXAMPLE_SPLITS_KEY,
'null'))
if not example_splits:
example_splits = ['eval']
logging.info("The 'example_splits' parameter is not set, using 'eval' "
'split.')
logging.info('Evaluating model.')
# TempPipInstallContext is needed here so that subprocesses (which
# may be created by the Beam multi-process DirectRunner) can find the
# needed dependencies.
# TODO(b/187122662): Move this to the ExecutorOperator or Launcher.
with udf_utils.TempPipInstallContext(extra_pip_packages):
with self._make_beam_pipeline() as pipeline:
examples_list = []
tensor_adapter_config = None
# pylint: disable=expression-not-assigned
if tfma.is_batched_input(eval_shared_model, eval_config):
tfxio_factory = tfxio_utils.get_tfxio_factory_from_artifact(
examples=input_dict[standard_component_specs.EXAMPLES_KEY],
telemetry_descriptors=_TELEMETRY_DESCRIPTORS,
schema=schema,
raw_record_column_name=tfma_constants.ARROW_INPUT_COLUMN)
# TODO(b/161935932): refactor after TFXIO supports multiple patterns.
for split in example_splits:
split_uris = artifact_utils.get_split_uris(
input_dict[standard_component_specs.EXAMPLES_KEY], split)
for index in range(len(split_uris)):
split_uri = split_uris[index]
file_pattern = io_utils.all_files_pattern(split_uri)
tfxio = tfxio_factory(file_pattern)
data = (
pipeline
| f'ReadFromTFRecordToArrow[{split}][{index}]' >>
tfxio.BeamSource())
examples_list.append(data)
if schema is not None:
# Use last tfxio as TensorRepresentations and ArrowSchema are fixed.
tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
arrow_schema=tfxio.ArrowSchema(),
tensor_representations=tfxio.TensorRepresentations())
else:
for split in example_splits:
split_uris = artifact_utils.get_split_uris(
input_dict[standard_component_specs.EXAMPLES_KEY], split)
for index in range(len(split_uris)):
split_uri = split_uris[index]
file_pattern = io_utils.all_files_pattern(split_uri)
data = (
pipeline
| f'ReadFromTFRecord[{split}][{index}]' >>
beam.io.ReadFromTFRecord(file_pattern=file_pattern))
examples_list.append(data)
custom_extractors = udf_utils.try_get_fn(
exec_properties=exec_properties, fn_name='custom_extractors')
extractors = None
if custom_extractors:
extractors = custom_extractors(
eval_shared_model=eval_shared_model,
eval_config=eval_config,
tensor_adapter_config=tensor_adapter_config)
(examples_list | 'FlattenExamples' >> beam.Flatten()
| 'ExtractEvaluateAndWriteResults' >>
(tfma.ExtractEvaluateAndWriteResults(
eval_shared_model=models[0] if len(models) == 1 else models,
eval_config=eval_config,
extractors=extractors,
output_path=output_uri,
slice_spec=slice_spec,
tensor_adapter_config=tensor_adapter_config)))
logging.info('Evaluation complete. Results written to %s.', output_uri)
if not run_validation:
# TODO(jinhuang): delete the BLESSING_KEY from output_dict when supported.
logging.info('No threshold configured, will not validate model.')
return
# Set up blessing artifact
blessing = artifact_utils.get_single_instance(
output_dict[standard_component_specs.BLESSING_KEY])
blessing.set_string_custom_property(
constants.ARTIFACT_PROPERTY_CURRENT_MODEL_URI_KEY,
artifact_utils.get_single_uri(
input_dict[standard_component_specs.MODEL_KEY]))
blessing.set_int_custom_property(
constants.ARTIFACT_PROPERTY_CURRENT_MODEL_ID_KEY,
input_dict[standard_component_specs.MODEL_KEY][0].id)
if input_dict.get(standard_component_specs.BASELINE_MODEL_KEY):
baseline_model = input_dict[
standard_component_specs.BASELINE_MODEL_KEY][0]
blessing.set_string_custom_property(
constants.ARTIFACT_PROPERTY_BASELINE_MODEL_URI_KEY,
baseline_model.uri)
blessing.set_int_custom_property(
constants.ARTIFACT_PROPERTY_BASELINE_MODEL_ID_KEY, baseline_model.id)
if 'current_component_id' in exec_properties:
blessing.set_string_custom_property(
'component_id', exec_properties['current_component_id'])
# Check validation result and write BLESSED file accordingly.
logging.info('Checking validation results.')
validation_result = tfma.load_validation_result(output_uri)
if validation_result.validation_ok:
io_utils.write_string_file(
os.path.join(blessing.uri, constants.BLESSED_FILE_NAME), '')
blessing.set_int_custom_property(constants.ARTIFACT_PROPERTY_BLESSED_KEY,
constants.BLESSED_VALUE)
else:
io_utils.write_string_file(
os.path.join(blessing.uri, constants.NOT_BLESSED_FILE_NAME), '')
blessing.set_int_custom_property(constants.ARTIFACT_PROPERTY_BLESSED_KEY,
constants.NOT_BLESSED_VALUE)
logging.info('Blessing result %s written to %s.',
validation_result.validation_ok, blessing.uri)
def testPredictExtractorWithRegressionModel(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.BATCHED_INPUT_KEY)
tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
arrow_schema=tfx_io.ArrowSchema(),
tensor_representations=tfx_io.TensorRepresentations())
input_extractor = batched_input_extractor.BatchedInputExtractor(
eval_config)
predict_extractor = batched_predict_extractor_v2.BatchedPredictExtractor(
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()
| input_extractor.stage_name >> input_extractor.ptransform
| predict_extractor.stage_name >> predict_extractor.ptransform)
# pylint: enable=no-value-for-parameter
def check_result(got):
try:
self.assertLen(got, 1)
self.assertIn(constants.BATCHED_PREDICTIONS_KEY, got[0])
expected_preds = [0.2, 0.8, 0.5]
self.assertAlmostEqual(
got[0][constants.BATCHED_PREDICTIONS_KEY],
expected_preds)
except AssertionError as err:
raise util.BeamAssertException(err)
util.assert_that(result, check_result, label='result')
def testWriteValidationResults(self):
model_dir, baseline_dir = self._getExportDir(), self._getBaselineDir()
eval_shared_model = self._build_keras_model(model_dir, mul=0)
baseline_eval_shared_model = self._build_keras_model(baseline_dir,
mul=1)
validations_file = os.path.join(self._getTempDir(),
constants.VALIDATIONS_KEY)
schema = text_format.Parse(
"""
tensor_representation_group {
key: ""
value {
tensor_representation {
key: "input"
value {
dense_tensor {
column_name: "input"
shape { dim { size: 1 } }
}
}
}
}
}
feature {
name: "input"
type: FLOAT
}
feature {
name: "label"
type: FLOAT
}
feature {
name: "example_weight"
type: FLOAT
}
feature {
name: "extra_feature"
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())
examples = [
self._makeExample(input=0.0,
label=1.0,
example_weight=1.0,
extra_feature='non_model_feature'),
self._makeExample(input=1.0,
label=0.0,
example_weight=0.5,
extra_feature='non_model_feature'),
]
eval_config = config.EvalConfig(
model_specs=[
config.ModelSpec(name='candidate',
label_key='label',
example_weight_key='example_weight'),
config.ModelSpec(name='baseline',
label_key='label',
example_weight_key='example_weight',
is_baseline=True)
],
slicing_specs=[config.SlicingSpec()],
metrics_specs=[
config.MetricsSpec(
metrics=[
config.MetricConfig(
class_name='WeightedExampleCount',
# 1.5 < 1, NOT OK.
threshold=config.MetricThreshold(
value_threshold=config.GenericValueThreshold(
upper_bound={'value': 1}))),
config.MetricConfig(
class_name='ExampleCount',
# 2 > 10, NOT OK.
threshold=config.MetricThreshold(
value_threshold=config.GenericValueThreshold(
lower_bound={'value': 10}))),
config.MetricConfig(
class_name='MeanLabel',
# 0 > 0 and 0 > 0%?: NOT OK.
threshold=config.MetricThreshold(
change_threshold=config.GenericChangeThreshold(
direction=config.MetricDirection.
HIGHER_IS_BETTER,
relative={'value': 0},
absolute={'value': 0}))),
config.MetricConfig(
# MeanPrediction = (0+0)/(1+0.5) = 0
class_name='MeanPrediction',
# -.01 < 0 < .01, OK.
# Diff% = -.333/.333 = -100% < -99%, OK.
# Diff = 0 - .333 = -.333 < 0, OK.
threshold=config.MetricThreshold(
value_threshold=config.GenericValueThreshold(
upper_bound={'value': .01},
lower_bound={'value': -.01}),
change_threshold=config.GenericChangeThreshold(
direction=config.MetricDirection.
LOWER_IS_BETTER,
relative={'value': -.99},
absolute={'value': 0})))
],
model_names=['candidate', 'baseline']),
],
options=config.Options(
disabled_outputs={'values': ['eval_config.json']}),
)
slice_spec = [
slicer.SingleSliceSpec(spec=s) for s in eval_config.slicing_specs
]
eval_shared_models = {
'candidate': eval_shared_model,
'baseline': baseline_eval_shared_model
}
extractors = [
batched_input_extractor.BatchedInputExtractor(eval_config),
batched_predict_extractor_v2.BatchedPredictExtractor(
eval_shared_model=eval_shared_models,
eval_config=eval_config,
tensor_adapter_config=tensor_adapter_config),
unbatch_extractor.UnbatchExtractor(),
slice_key_extractor.SliceKeyExtractor(slice_spec=slice_spec)
]
evaluators = [
metrics_and_plots_evaluator_v2.MetricsAndPlotsEvaluator(
eval_config=eval_config, eval_shared_model=eval_shared_models)
]
output_paths = {
constants.VALIDATIONS_KEY: validations_file,
}
writers = [
metrics_plots_and_validations_writer.
MetricsPlotsAndValidationsWriter(output_paths,
add_metrics_callbacks=[])
]
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()
|
'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts()
| 'ExtractEvaluate' >> model_eval_lib.ExtractAndEvaluate(
extractors=extractors, evaluators=evaluators)
|
'WriteResults' >> model_eval_lib.WriteResults(writers=writers))
# pylint: enable=no-value-for-parameter
validation_result = model_eval_lib.load_validation_result(
os.path.dirname(validations_file))
expected_validations = [
text_format.Parse(
"""
metric_key {
name: "weighted_example_count"
model_name: "candidate"
}
metric_threshold {
value_threshold {
upper_bound {
value: 1.0
}
}
}
metric_value {
double_value {
value: 1.5
}
}
""", validation_result_pb2.ValidationFailure()),
text_format.Parse(
"""
metric_key {
name: "example_count"
}
metric_threshold {
value_threshold {
lower_bound {
value: 10.0
}
}
}
metric_value {
double_value {
value: 2.0
}
}
""", validation_result_pb2.ValidationFailure()),
text_format.Parse(
"""
metric_key {
name: "mean_label"
model_name: "candidate"
is_diff: true
}
metric_threshold {
change_threshold {
absolute {
value: 0.0
}
relative {
value: 0.0
}
direction: HIGHER_IS_BETTER
}
}
metric_value {
double_value {
value: 0.0
}
}
""", validation_result_pb2.ValidationFailure()),
]
self.assertFalse(validation_result.validation_ok)
self.assertLen(validation_result.metric_validations_per_slice, 1)
self.assertCountEqual(
expected_validations,
validation_result.metric_validations_per_slice[0].failures)
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 testPredictionsExtractorWithMultiClassModel(self):
temp_export_dir = self._getExportDir()
export_dir, _ = dnn_classifier.simple_dnn_classifier(
temp_export_dir, None, n_classes=3)
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: "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)
prediction_extractor = predictions_extractor.PredictionsExtractor(
eval_config=eval_config,
eval_shared_model=eval_shared_model,
tensor_adapter_config=tensor_adapter_config)
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 in item[constants.PREDICTIONS_KEY]:
for pred_key in ('probabilities', 'all_classes'):
self.assertIn(pred_key, pred)
except AssertionError as err:
raise util.BeamAssertException(err)
util.assert_that(result, check_result, label='result')
def Do(self, input_dict: Dict[Text, List[types.Artifact]],
output_dict: Dict[Text, List[types.Artifact]],
exec_properties: Dict[Text, Any]) -> None:
"""Runs a batch job to evaluate the eval_model against the given input.
Args:
input_dict: Input dict from input key to a list of Artifacts.
- model_exports: exported model.
- examples: examples for eval the model.
output_dict: Output dict from output key to a list of Artifacts.
- output: model evaluation results.
exec_properties: A dict of execution properties.
- eval_config: JSON string of tfma.EvalConfig.
- feature_slicing_spec: JSON string of evaluator_pb2.FeatureSlicingSpec
instance, providing the way to slice the data. Deprecated, use
eval_config.slicing_specs instead.
Returns:
None
"""
if constants.EXAMPLES_KEY not in input_dict:
raise ValueError('EXAMPLES_KEY is missing from input dict.')
if constants.MODEL_KEY not in input_dict:
raise ValueError('MODEL_KEY is missing from input dict.')
if constants.EVALUATION_KEY not in output_dict:
raise ValueError('EVALUATION_KEY is missing from output dict.')
if len(input_dict[constants.MODEL_KEY]) > 1:
raise ValueError(
'There can be only one candidate model, there are {}.'.format(
len(input_dict[constants.MODEL_KEY])))
if constants.BASELINE_MODEL_KEY in input_dict and len(
input_dict[constants.BASELINE_MODEL_KEY]) > 1:
raise ValueError(
'There can be only one baseline model, there are {}.'.format(
len(input_dict[constants.BASELINE_MODEL_KEY])))
self._log_startup(input_dict, output_dict, exec_properties)
# Add fairness indicator metric callback if necessary.
fairness_indicator_thresholds = exec_properties.get(
'fairness_indicator_thresholds', None)
add_metrics_callbacks = None
if fairness_indicator_thresholds:
# Need to import the following module so that the fairness indicator
# post-export metric is registered.
import tensorflow_model_analysis.addons.fairness.post_export_metrics.fairness_indicators # pylint: disable=g-import-not-at-top, unused-variable
add_metrics_callbacks = [
tfma.post_export_metrics.fairness_indicators( # pytype: disable=module-attr
thresholds=fairness_indicator_thresholds),
]
output_uri = artifact_utils.get_single_uri(
output_dict[constants.EVALUATION_KEY])
run_validation = False
models = []
if 'eval_config' in exec_properties and exec_properties['eval_config']:
slice_spec = None
has_baseline = bool(input_dict.get(constants.BASELINE_MODEL_KEY))
eval_config = tfma.EvalConfig()
json_format.Parse(exec_properties['eval_config'], eval_config)
eval_config = tfma.update_eval_config_with_defaults(
eval_config,
maybe_add_baseline=has_baseline,
maybe_remove_baseline=not has_baseline)
tfma.verify_eval_config(eval_config)
# Do not validate model when there is no thresholds configured. This is to
# avoid accidentally blessing models when users forget to set thresholds.
run_validation = bool(
tfma.metrics.metric_thresholds_from_metrics_specs(
eval_config.metrics_specs))
if len(eval_config.model_specs) > 2:
raise ValueError(
"""Cannot support more than two models. There are {} models in this
eval_config.""".format(len(eval_config.model_specs)))
# Extract model artifacts.
for model_spec in eval_config.model_specs:
if model_spec.is_baseline:
model_uri = artifact_utils.get_single_uri(
input_dict[constants.BASELINE_MODEL_KEY])
else:
model_uri = artifact_utils.get_single_uri(
input_dict[constants.MODEL_KEY])
if tfma.get_model_type(model_spec) == tfma.TF_ESTIMATOR:
model_path = path_utils.eval_model_path(model_uri)
else:
model_path = path_utils.serving_model_path(model_uri)
absl.logging.info('Using {} as {} model.'.format(
model_path, model_spec.name))
models.append(
tfma.default_eval_shared_model(
model_name=model_spec.name,
eval_saved_model_path=model_path,
add_metrics_callbacks=add_metrics_callbacks,
eval_config=eval_config))
else:
eval_config = None
assert ('feature_slicing_spec' in exec_properties
and exec_properties['feature_slicing_spec']
), 'both eval_config and feature_slicing_spec are unset.'
feature_slicing_spec = evaluator_pb2.FeatureSlicingSpec()
json_format.Parse(exec_properties['feature_slicing_spec'],
feature_slicing_spec)
slice_spec = self._get_slice_spec_from_feature_slicing_spec(
feature_slicing_spec)
model_uri = artifact_utils.get_single_uri(
input_dict[constants.MODEL_KEY])
model_path = path_utils.eval_model_path(model_uri)
absl.logging.info('Using {} for model eval.'.format(model_path))
models.append(
tfma.default_eval_shared_model(
eval_saved_model_path=model_path,
add_metrics_callbacks=add_metrics_callbacks))
file_pattern = io_utils.all_files_pattern(
artifact_utils.get_split_uri(input_dict[constants.EXAMPLES_KEY],
'eval'))
eval_shared_model = models[0] if len(models) == 1 else models
schema = None
if constants.SCHEMA_KEY in input_dict:
schema = io_utils.SchemaReader().read(
io_utils.get_only_uri_in_dir(
artifact_utils.get_single_uri(
input_dict[constants.SCHEMA_KEY])))
absl.logging.info('Evaluating model.')
with self._make_beam_pipeline() as pipeline:
# pylint: disable=expression-not-assigned
if _USE_TFXIO:
tensor_adapter_config = None
if tfma.is_batched_input(eval_shared_model, eval_config):
tfxio = tf_example_record.TFExampleRecord(
file_pattern=file_pattern,
schema=schema,
raw_record_column_name=tfma.BATCHED_INPUT_KEY)
if schema is not None:
tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
arrow_schema=tfxio.ArrowSchema(),
tensor_representations=tfxio.TensorRepresentations(
))
data = pipeline | 'ReadFromTFRecordToArrow' >> tfxio.BeamSource(
)
else:
data = pipeline | 'ReadFromTFRecord' >> beam.io.ReadFromTFRecord(
file_pattern=file_pattern)
(data
| 'ExtractEvaluateAndWriteResults' >>
tfma.ExtractEvaluateAndWriteResults(
eval_shared_model=models[0]
if len(models) == 1 else models,
eval_config=eval_config,
output_path=output_uri,
slice_spec=slice_spec,
tensor_adapter_config=tensor_adapter_config))
else:
data = pipeline | 'ReadFromTFRecord' >> beam.io.ReadFromTFRecord(
file_pattern=file_pattern)
(data
| 'ExtractEvaluateAndWriteResults' >>
tfma.ExtractEvaluateAndWriteResults(
eval_shared_model=models[0]
if len(models) == 1 else models,
eval_config=eval_config,
output_path=output_uri,
slice_spec=slice_spec))
absl.logging.info(
'Evaluation complete. Results written to {}.'.format(output_uri))
if not run_validation:
# TODO(jinhuang): delete the BLESSING_KEY from output_dict when supported.
absl.logging.info(
'No threshold configured, will not validate model.')
return
# Set up blessing artifact
blessing = artifact_utils.get_single_instance(
output_dict[constants.BLESSING_KEY])
blessing.set_string_custom_property(
constants.ARTIFACT_PROPERTY_CURRENT_MODEL_URI_KEY,
artifact_utils.get_single_uri(input_dict[constants.MODEL_KEY]))
blessing.set_int_custom_property(
constants.ARTIFACT_PROPERTY_CURRENT_MODEL_ID_KEY,
input_dict[constants.MODEL_KEY][0].id)
if input_dict.get(constants.BASELINE_MODEL_KEY):
baseline_model = input_dict[constants.BASELINE_MODEL_KEY][0]
blessing.set_string_custom_property(
constants.ARTIFACT_PROPERTY_BASELINE_MODEL_URI_KEY,
baseline_model.uri)
blessing.set_int_custom_property(
constants.ARTIFACT_PROPERTY_BASELINE_MODEL_ID_KEY,
baseline_model.id)
if 'current_component_id' in exec_properties:
blessing.set_string_custom_property(
'component_id', exec_properties['current_component_id'])
# Check validation result and write BLESSED file accordingly.
absl.logging.info('Checking validation results.')
validation_result = tfma.load_validation_result(output_uri)
if validation_result.validation_ok:
io_utils.write_string_file(
os.path.join(blessing.uri, constants.BLESSED_FILE_NAME), '')
blessing.set_int_custom_property(
constants.ARTIFACT_PROPERTY_BLESSED_KEY,
constants.BLESSED_VALUE)
else:
io_utils.write_string_file(
os.path.join(blessing.uri, constants.NOT_BLESSED_FILE_NAME),
'')
blessing.set_int_custom_property(
constants.ARTIFACT_PROPERTY_BLESSED_KEY,
constants.NOT_BLESSED_VALUE)
absl.logging.info('Blessing result {} written to {}.'.format(
validation_result.validation_ok, blessing.uri))
def testPredictExtractorWithMultiModels(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.BATCHED_INPUT_KEY)
tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
arrow_schema=tfx_io.ArrowSchema(),
tensor_representations=tfx_io.TensorRepresentations())
input_extractor = batched_input_extractor.BatchedInputExtractor(
eval_config)
predict_extractor = batched_predict_extractor_v2.BatchedPredictExtractor(
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()
| input_extractor.stage_name >> input_extractor.ptransform
| predict_extractor.stage_name >> predict_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.BATCHED_PREDICTIONS_KEY, item)
for pred in item[constants.BATCHED_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 _batch_reducible_process(
self, batched_extract: types.Extracts) -> List[types.Extracts]:
def maybe_expand_dims(arr):
if not hasattr(arr, 'shape') or not arr.shape:
return np.expand_dims(arr, axis=0)
else:
return arr
def to_dense(t):
return tf.sparse.to_dense(t) if isinstance(t,
tf.SparseTensor) else t
result = copy.copy(batched_extract)
record_batch = batched_extract[constants.ARROW_RECORD_BATCH_KEY]
serialized_examples = batched_extract[constants.INPUT_KEY]
for extracts_key in self._signature_names.keys():
if extracts_key not in result or not result[extracts_key]:
result[extracts_key] = [None] * record_batch.num_rows
for model_name, model in self._loaded_models.items():
for extracts_key, signature_names in self._signature_names.items():
for signature_name in (signature_names[model_name]
or self._default_signature_names):
required = bool(signature_names[model_name])
input_specs = get_input_specs(model, signature_name,
required) or {}
inputs = None
# If input_specs exist then try to filter the inputs by the input
# names (unlike estimators, keras does not accept unknown inputs).
if input_specs:
adapter = self._tensor_adapter
if (not adapter and set(input_specs.keys()) <= set(
record_batch.schema.names)):
# Create adapter based on input_specs
tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
arrow_schema=record_batch.schema,
tensor_representations=
input_specs_to_tensor_representations(
input_specs))
adapter = tensor_adapter.TensorAdapter(
tensor_adapter_config)
# Avoid getting the tensors if we appear to be feeding serialized
# examples to the callable.
if adapter and not (
len(input_specs) == 1
and next(iter(input_specs.values())).dtype
== tf.string and find_input_name_in_features(
set(adapter.TypeSpecs().keys()),
next(iter(input_specs.keys()))) is None):
# TODO(b/172376802): Update to pass input specs to ToBatchTensors.
inputs = filter_by_input_names(
adapter.ToBatchTensors(record_batch),
list(input_specs.keys()))
if not inputs:
# Assume serialized examples
assert serialized_examples is not None, 'Raw examples not found.'
inputs = serialized_examples
# If a signature name was not provided, default to using the serving
# signature since parsing normally will be done outside model.
if not signature_name:
signature_name = get_default_signature_name(model)
signature = get_callable(model, signature_name, required)
if signature is None:
if not required:
continue
raise ValueError(
'Unable to find %s function needed to update %s' %
(signature_name, extracts_key))
if isinstance(inputs, dict):
if hasattr(signature, 'structured_input_signature'):
outputs = signature(**inputs)
else:
outputs = signature(inputs)
else:
outputs = signature(
tf.constant(inputs, dtype=tf.string))
for i in range(record_batch.num_rows):
if isinstance(outputs, dict):
output = {
k: maybe_expand_dims(to_dense(v)[i].numpy())
for k, v in outputs.items()
}
else:
output = {
signature_name:
maybe_expand_dims(
np.asarray(to_dense(outputs))[i])
}
if result[extracts_key][i] is None:
result[extracts_key][i] = collections.defaultdict(
dict)
result[extracts_key][i][model_name].update(output) # pytype: disable=unsupported-operands
for i in range(len(result[extracts_key])):
# PyType doesn't recognize isinstance(..., dict).
# pytype: disable=attribute-error,unsupported-operands
if isinstance(result[extracts_key][i], dict):
for model_name, output in result[extracts_key][i].items():
if not self._prefer_dict_outputs and len(output) == 1:
result[extracts_key][i][model_name] = list(
output.values())[0]
# If only one model, the output is stored without using a dict
if len(self._eval_config.model_specs) == 1:
result[extracts_key][i] = list(
result[extracts_key][i].values())[0]
# pytype: enable=attribute-error,unsupported-operands
return [result]
def testPredictExtractorWithSequentialKerasModel(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.BATCHED_INPUT_KEY)
tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
arrow_schema=tfx_io.ArrowSchema(),
tensor_representations=tfx_io.TensorRepresentations())
input_extractor = batched_input_extractor.BatchedInputExtractor(
eval_config)
predict_extractor = batched_predict_extractor_v2.BatchedPredictExtractor(
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()
| input_extractor.stage_name >> input_extractor.ptransform
| predict_extractor.stage_name >> predict_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.BATCHED_PREDICTIONS_KEY, item)
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.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 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.ModelSpec(name=model_name))
eval_config = config.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 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')
