def testBatchSizeLimit(self):
temp_eval_export_dir = self._getEvalExportDir()
_, eval_export_dir = batch_size_limited_classifier.simple_batch_size_limited_classifier(
None, temp_eval_export_dir)
eval_shared_model = model_eval_lib.default_eval_shared_model(
eval_saved_model_path=eval_export_dir)
with beam.Pipeline() as pipeline:
examples = [
self._makeExample(classes='first', scores=0.0, labels='third'),
self._makeExample(classes='first', scores=0.0, labels='third'),
self._makeExample(classes='first', scores=0.0, labels='third'),
self._makeExample(classes='first', scores=0.0, labels='third'),
]
serialized_examples = [e.SerializeToString() for e in examples]
predict_extracts = (
pipeline
| beam.Create(serialized_examples, reshuffle=False)
# Our diagnostic outputs, pass types.Extracts throughout, however our
# aggregating functions do not use this interface.
| beam.Map(lambda x: {constants.INPUT_KEY: x})
| 'Predict' >> predict_extractor._TFMAPredict(
eval_shared_models={'': eval_shared_model}))
def check_result(got):
self.assertLen(got, 4)
for item in got:
self.assertIn(constants.PREDICTIONS_KEY, item)
util.assert_that(predict_extracts, check_result)
def testBatchSizeLimit(self):
temp_export_dir = self._getExportDir()
_, export_dir = batch_size_limited_classifier.simple_batch_size_limited_classifier(
None, temp_export_dir)
eval_shared_model = self.createTestEvalSharedModel(
eval_saved_model_path=export_dir, tags=[tf.saved_model.SERVING])
eval_config = config.EvalConfig(model_specs=[config.ModelSpec()])
predict_extractor = predict_extractor_v2.PredictExtractor(
eval_config=eval_config, eval_shared_model=eval_shared_model)
with beam.Pipeline() as pipeline:
examples = [
self._makeExample(classes='first', scores=0.0, labels='third'),
self._makeExample(classes='first', scores=0.0, labels='third'),
self._makeExample(classes='first', scores=0.0, labels='third'),
self._makeExample(classes='first', scores=0.0, labels='third'),
]
predict_extracts = (
pipeline
| 'Create' >> beam.Create([e.SerializeToString() for e in examples],
reshuffle=False)
| 'FeaturesToExtracts' >> model_eval_lib.InputsToExtracts()
| 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.PREDICTIONS_KEY, item)
except AssertionError as err:
raise util.BeamAssertException(err)
util.assert_that(predict_extracts, check_result, label='result')
def testBatchSizeLimit(self):
temp_export_dir = self._getExportDir()
_, export_dir = batch_size_limited_classifier.simple_batch_size_limited_classifier(
None, temp_export_dir)
eval_shared_model = self.createTestEvalSharedModel(
eval_saved_model_path=export_dir, tags=[tf.saved_model.SERVING])
eval_config = config.EvalConfig(model_specs=[config.ModelSpec()])
schema = text_format.Parse(
"""
feature {
name: "classes"
type: BYTES
}
feature {
name: "scores"
type: FLOAT
}
feature {
name: "labels"
type: BYTES
}
""", schema_pb2.Schema())
tfx_io = test_util.InMemoryTFExampleRecord(
schema=schema, raw_record_column_name=constants.ARROW_INPUT_COLUMN)
tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
arrow_schema=tfx_io.ArrowSchema(),
tensor_representations=tfx_io.TensorRepresentations())
feature_extractor = features_extractor.FeaturesExtractor(eval_config)
prediction_extractor = predictions_extractor.PredictionsExtractor(
eval_config=eval_config,
eval_shared_model=eval_shared_model,
tensor_adapter_config=tensor_adapter_config)
examples = []
for _ in range(4):
examples.append(
self._makeExample(classes='first', scores=0.0, labels='third'))
with beam.Pipeline() as pipeline:
predict_extracts = (
pipeline
| 'Create' >> beam.Create([e.SerializeToString() for e in examples],
reshuffle=False)
| 'BatchExamples' >> tfx_io.BeamSource(batch_size=1)
| 'InputsToExtracts' >> model_eval_lib.BatchedInputsToExtracts()
| feature_extractor.stage_name >> feature_extractor.ptransform
| prediction_extractor.stage_name >> prediction_extractor.ptransform)
def check_result(got):
try:
self.assertLen(got, 4)
# We can't verify the actual predictions, but we can verify the keys.
for item in got:
self.assertIn(constants.PREDICTIONS_KEY, item)
except AssertionError as err:
raise util.BeamAssertException(err)
util.assert_that(predict_extracts, check_result, label='result')