def testConvertNamedSignatureNonGenericToSignatureDef(self):
signatures_proto = manifest_pb2.Signatures()
regression_signature = manifest_pb2.RegressionSignature()
signatures_proto.named_signatures[
signature_constants.PREDICT_INPUTS].regression_signature.CopyFrom(
regression_signature)
with self.assertRaises(RuntimeError):
_ = bundle_shim._convert_named_signatures_to_signature_def(
signatures_proto)
signatures_proto = manifest_pb2.Signatures()
classification_signature = manifest_pb2.ClassificationSignature()
signatures_proto.named_signatures[
signature_constants.PREDICT_INPUTS].classification_signature.CopyFrom(
classification_signature)
with self.assertRaises(RuntimeError):
_ = bundle_shim._convert_named_signatures_to_signature_def(
signatures_proto)
def testConvertDefaultSignatureClassificationToSignatureDef(self):
signatures_proto = manifest_pb2.Signatures()
classification_signature = manifest_pb2.ClassificationSignature()
classification_signature.input.CopyFrom(
manifest_pb2.TensorBinding(
tensor_name=signature_constants.CLASSIFY_INPUTS))
classification_signature.classes.CopyFrom(
manifest_pb2.TensorBinding(
tensor_name=signature_constants.CLASSIFY_OUTPUT_CLASSES))
classification_signature.scores.CopyFrom(
manifest_pb2.TensorBinding(
tensor_name=signature_constants.CLASSIFY_OUTPUT_SCORES))
signatures_proto.default_signature.classification_signature.CopyFrom(
classification_signature)
signatures_proto.default_signature.classification_signature.CopyFrom(
classification_signature)
signature_def = bundle_shim._convert_default_signature_to_signature_def(
signatures_proto)
# Validate classification signature correctly copied over.
self.assertEqual(signature_def.method_name,
signature_constants.CLASSIFY_METHOD_NAME)
self.assertEqual(len(signature_def.inputs), 1)
self.assertEqual(len(signature_def.outputs), 2)
self.assertProtoEquals(
signature_def.inputs[signature_constants.CLASSIFY_INPUTS],
meta_graph_pb2.TensorInfo(
name=signature_constants.CLASSIFY_INPUTS))
self.assertProtoEquals(
signature_def.outputs[signature_constants.CLASSIFY_OUTPUT_SCORES],
meta_graph_pb2.TensorInfo(
name=signature_constants.CLASSIFY_OUTPUT_SCORES))
self.assertProtoEquals(
signature_def.outputs[signature_constants.CLASSIFY_OUTPUT_CLASSES],
meta_graph_pb2.TensorInfo(
name=signature_constants.CLASSIFY_OUTPUT_CLASSES))