def test_build_all_signature_defs_serving_only(self): receiver_tensor = {"input": array_ops.placeholder(dtypes.string)} output_1 = constant_op.constant([1.]) export_outputs = { signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: export_output.PredictOutput(outputs=output_1), "train": export_output.TrainOutput(loss=output_1), } signature_defs = export.build_all_signature_defs( receiver_tensor, export_outputs) expected_signature_defs = { "serving_default": signature_def_utils.predict_signature_def( receiver_tensor, {"output": output_1}) } self.assertDictEqual(expected_signature_defs, signature_defs) signature_defs = export.build_all_signature_defs( receiver_tensor, export_outputs, serving_only=False) expected_signature_defs.update({ "train": signature_def_utils.supervised_train_signature_def( receiver_tensor, loss={"loss": output_1}) }) self.assertDictEqual(expected_signature_defs, signature_defs)
def test_build_all_signature_defs_export_outputs_required(self): receiver_tensor = constant_op.constant(["11"]) with self.assertRaises(ValueError) as e: export.build_all_signature_defs(receiver_tensor, None) self.assertTrue(str(e.exception).startswith( "export_outputs must be a dict"))
def test_build_all_signature_defs_without_receiver_alternatives(self): receiver_tensor = array_ops.placeholder(dtypes.string) output_1 = constant_op.constant([1.]) output_2 = constant_op.constant(["2"]) output_3 = constant_op.constant(["3"]) export_outputs = { signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: export_output.RegressionOutput(value=output_1), "head-2": export_output.ClassificationOutput(classes=output_2), "head-3": export_output.PredictOutput(outputs={ "some_output_3": output_3 }), } signature_defs = export.build_all_signature_defs( receiver_tensor, export_outputs) expected_signature_defs = { "serving_default": signature_def_utils.regression_signature_def(receiver_tensor, output_1), "head-2": signature_def_utils.classification_signature_def(receiver_tensor, output_2, None), "head-3": signature_def_utils.predict_signature_def({ "input": receiver_tensor }, {"some_output_3": output_3}) } self.assertDictEqual(expected_signature_defs, signature_defs)
def test_build_all_signature_defs_with_single_alternatives(self): receiver_tensor = array_ops.placeholder(dtypes.string) receiver_tensors_alternative_1 = array_ops.placeholder(dtypes.int64) receiver_tensors_alternative_2 = array_ops.sparse_placeholder( dtypes.float32) # Note we are passing single Tensors as values of # receiver_tensors_alternatives, where normally that is a dict. # In this case a dict will be created using the default receiver tensor # name "input". receiver_tensors_alternatives = {"other1": receiver_tensors_alternative_1, "other2": receiver_tensors_alternative_2} output_1 = constant_op.constant([1.]) output_2 = constant_op.constant(["2"]) output_3 = constant_op.constant(["3"]) export_outputs = { signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: export_output.RegressionOutput(value=output_1), "head-2": export_output.ClassificationOutput(classes=output_2), "head-3": export_output.PredictOutput(outputs={ "some_output_3": output_3 }), } signature_defs = export.build_all_signature_defs( receiver_tensor, export_outputs, receiver_tensors_alternatives) expected_signature_defs = { "serving_default": signature_def_utils.regression_signature_def( receiver_tensor, output_1), "head-2": signature_def_utils.classification_signature_def( receiver_tensor, output_2, None), "head-3": signature_def_utils.predict_signature_def( {"input": receiver_tensor}, {"some_output_3": output_3}), "other1:head-3": signature_def_utils.predict_signature_def( {"input": receiver_tensors_alternative_1}, {"some_output_3": output_3}), "other2:head-3": signature_def_utils.predict_signature_def( {"input": receiver_tensors_alternative_2}, {"some_output_3": output_3}) # Note that the alternatives 'other:serving_default' and 'other:head-2' # are invalid, because regession and classification signatures must take # a single string input. Here we verify that these invalid signatures # are not included in the export. } self.assertDictEqual(expected_signature_defs, signature_defs)
def _create_signature_def_map(model, mode): """Creates a SignatureDef map from a Keras model.""" inputs_dict = {name: x for name, x in zip(model.input_names, model.inputs)} if model.optimizer: targets_dict = { x.name.split(':')[0]: x for x in model.targets if x is not None } inputs_dict.update(targets_dict) outputs_dict = { name: x for name, x in zip(model.output_names, model.outputs) } metrics = estimator_keras_util._convert_keras_metrics_to_estimator(model) # Add metric variables to the `LOCAL_VARIABLES` collection. Metric variables # are by default not added to any collections. We are doing this here, so # that metric variables get initialized. local_vars = set(ops.get_collection(ops.GraphKeys.LOCAL_VARIABLES)) vars_to_add = set() if metrics is not None: for key, value in six.iteritems(metrics): if isinstance(value, Metric): vars_to_add.update(value.variables) # Convert Metric instances to (value_tensor, update_op) tuple. metrics[key] = (value.result(), value.updates[0]) # Remove variables that are in the local variables collection already. vars_to_add = vars_to_add.difference(local_vars) for v in vars_to_add: ops.add_to_collection(ops.GraphKeys.LOCAL_VARIABLES, v) export_outputs = model_fn_lib.export_outputs_for_mode( mode, predictions=outputs_dict, loss=model.total_loss if model.optimizer else None, metrics=metrics) return export_helpers.build_all_signature_defs( inputs_dict, export_outputs=export_outputs, serving_only=(mode == model_fn_lib.ModeKeys.PREDICT))