def testConvertDefaultSignatureRegressionToSignatureDef(self):
signatures_proto = manifest_pb2.Signatures()
regression_signature = manifest_pb2.RegressionSignature()
regression_signature.input.CopyFrom(
manifest_pb2.TensorBinding(
tensor_name=signature_constants.REGRESS_INPUTS))
regression_signature.output.CopyFrom(
manifest_pb2.TensorBinding(
tensor_name=signature_constants.REGRESS_OUTPUTS))
signatures_proto.default_signature.regression_signature.CopyFrom(
regression_signature)
signature_def = bundle_shim._convert_default_signature_to_signature_def(
signatures_proto)
# Validate regression signature correctly copied over.
self.assertEqual(signature_def.method_name,
signature_constants.REGRESS_METHOD_NAME)
self.assertEqual(len(signature_def.inputs), 1)
self.assertEqual(len(signature_def.outputs), 1)
self.assertProtoEquals(
signature_def.inputs[signature_constants.REGRESS_INPUTS],
meta_graph_pb2.TensorInfo(name=signature_constants.REGRESS_INPUTS))
self.assertProtoEquals(
signature_def.outputs[signature_constants.REGRESS_OUTPUTS],
meta_graph_pb2.TensorInfo(
name=signature_constants.REGRESS_OUTPUTS))
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)