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