def testTFlitePredictExtractorWithSingleOutputModel(
self, multi_model, multi_output, batch_examples, batch_inputs):
input1 = tf.keras.layers.Input(shape=(1, ), name='input1')
input2 = tf.keras.layers.Input(shape=(1, ), name='input2')
inputs = [input1, input2]
input_layer = tf.keras.layers.concatenate(inputs)
output_layers = {}
output_layers['output1'] = (tf.keras.layers.Dense(
1, activation=tf.nn.sigmoid, name='output1')(input_layer))
if multi_output:
output_layers['output2'] = (tf.keras.layers.Dense(
1, activation=tf.nn.sigmoid, name='output2')(input_layer))
model = tf.keras.models.Model(inputs, output_layers)
model.compile(optimizer=tf.keras.optimizers.Adam(lr=.001),
loss=tf.keras.losses.binary_crossentropy,
metrics=['accuracy'])
train_features = {'input1': [[0.0], [1.0]], 'input2': [[1.0], [0.0]]}
labels = {'output1': [[1], [0]]}
if multi_output:
labels['output2'] = [[1], [0]]
example_weights = {'output1': [1.0, 0.5]}
if multi_output:
example_weights['output2'] = [1.0, 0.5]
dataset = tf.data.Dataset.from_tensor_slices(
(train_features, labels, example_weights))
dataset = dataset.shuffle(buffer_size=1).repeat().batch(2)
model.fit(dataset, steps_per_epoch=1)
converter = tf.compat.v2.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
tflite_model_dir = tempfile.mkdtemp()
with tf.io.gfile.GFile(os.path.join(tflite_model_dir, 'tflite'),
'wb') as f:
f.write(tflite_model)
model_specs = [config.ModelSpec(name='model1')]
if multi_model:
model_specs.append(config.ModelSpec(name='model2'))
eval_config = config.EvalConfig(model_specs=model_specs)
eval_shared_model = self.createTestEvalSharedModel(
eval_saved_model_path=tflite_model_dir)
eval_shared_models = {'model1': eval_shared_model}
if multi_model:
eval_shared_models['model2'] = eval_shared_model
desired_batch_size = 2 if batch_examples else None
predictor = tflite_predict_extractor.TFLitePredictExtractor(
eval_config=eval_config,
eval_shared_model=eval_shared_models,
desired_batch_size=desired_batch_size)
predict_features = [
{
'input1': np.array([0.0], dtype=np.float32),
'input2': np.array([1.0], dtype=np.float32),
'non_model_feature':
np.array([0]), # should be ignored by model
},
{
'input1': np.array([1.0], dtype=np.float32),
'input2': np.array([0.0], dtype=np.float32),
'non_model_feature':
np.array([1]), # should be ignored by model
}
]
if batch_inputs:
predict_features = [{
k: np.expand_dims(v, 0)
for k, v in p.items()
} for p in predict_features]
with beam.Pipeline() as pipeline:
# pylint: disable=no-value-for-parameter
result = (
pipeline
| 'Create' >> beam.Create(predict_features, reshuffle=False)
| 'FeaturesToExtracts' >>
beam.Map(lambda x: {constants.FEATURES_KEY: x})
| predictor.stage_name >> predictor.ptransform)
# pylint: enable=no-value-for-parameter
def check_result(got):
try:
self.assertLen(got, 2)
# We can't verify the actual predictions, but we can verify the keys.
for item in got:
self.assertIn(constants.PREDICTIONS_KEY, item)
# TODO(dzats): TFLite seems to currently rename all outputs to
# Identity*. Update this test to check for output1 and output2
# when this changes.
if multi_model:
self.assertIn('model1',
item[constants.PREDICTIONS_KEY])
self.assertIn('model2',
item[constants.PREDICTIONS_KEY])
if multi_output:
self.assertIn(
'Identity',
item[constants.PREDICTIONS_KEY]['model1'])
self.assertIn(
'Identity_1',
item[constants.PREDICTIONS_KEY]['model1'])
elif multi_output:
self.assertIn('Identity',
item[constants.PREDICTIONS_KEY])
self.assertIn('Identity_1',
item[constants.PREDICTIONS_KEY])
except AssertionError as err:
raise util.BeamAssertException(err)
util.assert_that(result, check_result, label='result')
def testTFlitePredictExtractorWithKerasModel(self, multi_model,
multi_output):
input1 = tf.keras.layers.Input(shape=(1, ), name='input1')
input2 = tf.keras.layers.Input(shape=(1, ), name='input2')
inputs = [input1, input2]
input_layer = tf.keras.layers.concatenate(inputs)
output_layers = {}
output_layers['output1'] = (tf.keras.layers.Dense(
1, activation=tf.nn.sigmoid, name='output1')(input_layer))
if multi_output:
output_layers['output2'] = (tf.keras.layers.Dense(
1, activation=tf.nn.sigmoid, name='output2')(input_layer))
model = tf.keras.models.Model(inputs, output_layers)
model.compile(optimizer=tf.keras.optimizers.Adam(lr=.001),
loss=tf.keras.losses.binary_crossentropy,
metrics=['accuracy'])
train_features = {'input1': [[0.0], [1.0]], 'input2': [[1.0], [0.0]]}
labels = {'output1': [[1], [0]]}
if multi_output:
labels['output2'] = [[1], [0]]
example_weights = {'output1': [1.0, 0.5]}
if multi_output:
example_weights['output2'] = [1.0, 0.5]
dataset = tf.data.Dataset.from_tensor_slices(
(train_features, labels, example_weights))
dataset = dataset.shuffle(buffer_size=1).repeat().batch(2)
model.fit(dataset, steps_per_epoch=1)
converter = tf.compat.v2.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
tflite_model_dir = tempfile.mkdtemp()
with tf.io.gfile.GFile(os.path.join(tflite_model_dir, 'tflite'),
'wb') as f:
f.write(tflite_model)
model_specs = [config.ModelSpec(name='model1', model_type='tf_lite')]
if multi_model:
model_specs.append(
config.ModelSpec(name='model2', model_type='tf_lite'))
eval_config = config.EvalConfig(model_specs=model_specs)
eval_shared_models = [
self.createTestEvalSharedModel(
model_name='model1',
eval_saved_model_path=tflite_model_dir,
model_type='tf_lite')
]
if multi_model:
eval_shared_models.append(
self.createTestEvalSharedModel(
model_name='model2',
eval_saved_model_path=tflite_model_dir,
model_type='tf_lite'))
schema = text_format.Parse(
"""
feature {
name: "input1"
type: FLOAT
}
feature {
name: "input2"
type: FLOAT
}
feature {
name: "non_model_feature"
type: INT
}
""", schema_pb2.Schema())
tfx_io = test_util.InMemoryTFExampleRecord(
schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN)
feature_extractor = features_extractor.FeaturesExtractor(eval_config)
predictor = tflite_predict_extractor.TFLitePredictExtractor(
eval_config=eval_config, eval_shared_model=eval_shared_models)
examples = [
self._makeExample(input1=0.0, input2=1.0, non_model_feature=0),
self._makeExample(input1=1.0, input2=0.0, non_model_feature=1),
]
with beam.Pipeline() as pipeline:
# pylint: disable=no-value-for-parameter
result = (
pipeline
| 'Create' >> beam.Create(
[e.SerializeToString() for e in examples], reshuffle=False)
| 'BatchExamples' >> tfx_io.BeamSource(batch_size=2)
|
'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts()
| feature_extractor.stage_name >> feature_extractor.ptransform
| predictor.stage_name >> predictor.ptransform)
# pylint: enable=no-value-for-parameter
def check_result(got):
try:
self.assertLen(got, 1)
got = got[0]
self.assertIn(constants.PREDICTIONS_KEY, got)
self.assertLen(got[constants.PREDICTIONS_KEY], 2)
for item in got[constants.PREDICTIONS_KEY]:
if multi_model:
self.assertIn('model1', item)
self.assertIn('model2', item)
if multi_output:
self.assertIn('Identity', item['model1'])
self.assertIn('Identity_1', item['model1'])
elif multi_output:
self.assertIn('Identity', item)
self.assertIn('Identity_1', item)
except AssertionError as err:
raise util.BeamAssertException(err)
util.assert_that(result, check_result, label='result')
def default_extractors( # pylint: disable=invalid-name
eval_shared_model: Union[types.EvalSharedModel,
Dict[Text, types.EvalSharedModel]] = None,
eval_config: config.EvalConfig = None,
slice_spec: Optional[List[slicer.SingleSliceSpec]] = None,
desired_batch_size: Optional[int] = None,
materialize: Optional[bool] = True) -> List[extractor.Extractor]:
"""Returns the default extractors for use in ExtractAndEvaluate.
Args:
eval_shared_model: Shared model (single-model evaluation) or dict of shared
models keyed by model name (multi-model evaluation). Required unless the
predictions are provided alongside of the features (i.e. model-agnostic
evaluations).
eval_config: Eval config.
slice_spec: Deprecated (use EvalConfig).
desired_batch_size: Optional batch size for batching in Predict.
materialize: True to have extractors create materialized output.
Raises:
NotImplementedError: If eval_config contains mixed serving and eval models.
"""
if eval_config is not None:
eval_config = config.update_eval_config_with_defaults(eval_config)
slice_spec = [
slicer.SingleSliceSpec(spec=spec) for spec in eval_config.slicing_specs
]
if _is_legacy_eval(eval_shared_model, eval_config):
# Backwards compatibility for previous add_metrics_callbacks implementation.
return [
predict_extractor.PredictExtractor(
eval_shared_model, desired_batch_size, materialize=materialize),
slice_key_extractor.SliceKeyExtractor(
slice_spec, materialize=materialize)
]
elif eval_shared_model:
model_types = model_util.get_model_types(eval_config)
if not model_types.issubset(constants.VALID_MODEL_TYPES):
raise NotImplementedError(
'model type must be one of: {}. evalconfig={}'.format(
str(constants.VALID_MODEL_TYPES), eval_config))
if model_types == set([constants.TF_LITE]):
return [
input_extractor.InputExtractor(eval_config=eval_config),
tflite_predict_extractor.TFLitePredictExtractor(
eval_config=eval_config,
eval_shared_model=eval_shared_model,
desired_batch_size=desired_batch_size),
slice_key_extractor.SliceKeyExtractor(
slice_spec, materialize=materialize)
]
elif constants.TF_LITE in model_types:
raise NotImplementedError(
'support for mixing tf_lite and non-tf_lite models is not '
'implemented: eval_config={}'.format(eval_config))
elif (eval_config and all(s.signature_name == eval_constants.EVAL_TAG
for s in eval_config.model_specs)):
return [
predict_extractor.PredictExtractor(
eval_shared_model,
desired_batch_size,
materialize=materialize,
eval_config=eval_config),
slice_key_extractor.SliceKeyExtractor(
slice_spec, materialize=materialize)
]
elif (eval_config and any(s.signature_name == eval_constants.EVAL_TAG
for s in eval_config.model_specs)):
raise NotImplementedError(
'support for mixing eval and non-eval models is not implemented: '
'eval_config={}'.format(eval_config))
else:
return [
input_extractor.InputExtractor(eval_config=eval_config),
predict_extractor_v2.PredictExtractor(
eval_config=eval_config,
eval_shared_model=eval_shared_model,
desired_batch_size=desired_batch_size),
slice_key_extractor.SliceKeyExtractor(
slice_spec, materialize=materialize)
]
else:
return [
input_extractor.InputExtractor(eval_config=eval_config),
slice_key_extractor.SliceKeyExtractor(
slice_spec, materialize=materialize)
]
