diff options
| author |  Michael Case <mikecase@google.com> | 2018-02-07 14:36:00 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-02-07 14:39:49 -0800 |
| commit | d90054e7c0f41f4bab81df0548577a73b939a87a (patch) | |
| tree | a15aea686a9d3f305e316d2a6ada0859ad8170d1 /tensorflow/contrib/session_bundle | |
| parent | 8461760f9f6cde8ed97507484d2a879140141032 (diff) | |
Merge changes from github.
PiperOrigin-RevId: 184897758
Diffstat (limited to 'tensorflow/contrib/session_bundle')
| -rw-r--r-- | tensorflow/contrib/session_bundle/bundle_shim.py | 11 | ||||
| -rw-r--r-- | tensorflow/contrib/session_bundle/constants.py | 3 |
2 files changed, 10 insertions, 4 deletions
diff --git a/tensorflow/contrib/session_bundle/bundle_shim.py b/tensorflow/contrib/session_bundle/bundle_shim.py index 062c9cc680..1db97020a2 100644 --- a/tensorflow/contrib/session_bundle/bundle_shim.py +++ b/tensorflow/contrib/session_bundle/bundle_shim.py @@ -82,7 +82,8 @@ def _convert_default_signature_to_signature_def(signatures): """ default_signature = signatures.default_signature signature_def = meta_graph_pb2.SignatureDef() - if default_signature.WhichOneof("type") == "regression_signature": + if (default_signature.WhichOneof("type") == + legacy_constants.REGRESSION_SIGNATURE): regression_signature = default_signature.regression_signature signature_def.method_name = signature_constants.REGRESS_METHOD_NAME _add_input_to_signature_def(regression_signature.input.tensor_name, @@ -91,7 +92,8 @@ def _convert_default_signature_to_signature_def(signatures): _add_output_to_signature_def(regression_signature.output.tensor_name, signature_constants.REGRESS_OUTPUTS, signature_def) - elif default_signature.WhichOneof("type") == "classification_signature": + elif (default_signature.WhichOneof("type") == + legacy_constants.CLASSIFICATION_SIGNATURE): classification_signature = default_signature.classification_signature signature_def.method_name = signature_constants.CLASSIFY_METHOD_NAME _add_input_to_signature_def(classification_signature.input.tensor_name, @@ -132,8 +134,9 @@ def _convert_named_signatures_to_signature_def(signatures): signature_constants.PREDICT_OUTPUTS] # TODO(pdudnik): what if there are other signatures? Mimic cr/140900781 once # it is submitted. - if (input_signature.WhichOneof("type") != "generic_signature" or - output_signature.WhichOneof("type") != "generic_signature"): + if (input_signature.WhichOneof("type") != legacy_constants.GENERIC_SIGNATURE + or output_signature.WhichOneof("type") != + legacy_constants.GENERIC_SIGNATURE): raise RuntimeError("Named input and output signatures can only be " "up-converted if they are generic signature. " "Input signature type is %s, output signature type is " diff --git a/tensorflow/contrib/session_bundle/constants.py b/tensorflow/contrib/session_bundle/constants.py index 6ced73241a..e833baee79 100644 --- a/tensorflow/contrib/session_bundle/constants.py +++ b/tensorflow/contrib/session_bundle/constants.py @@ -32,3 +32,6 @@ INIT_OP_KEY = "serving_init_op" SIGNATURES_KEY = "serving_signatures" ASSETS_KEY = "serving_assets" GRAPH_KEY = "serving_graph" +REGRESSION_SIGNATURE = "regression_signature" +CLASSIFICATION_SIGNATURE = "classification_signature" +GENERIC_SIGNATURE = "generic_signature" |