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) ]