405 lines
14 KiB
Python
405 lines
14 KiB
Python
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
# LINT.IfChange
|
|
"""Classes for different types of export output."""
|
|
|
|
import abc
|
|
|
|
from tensorflow.python.framework import constant_op
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.framework import tensor_util
|
|
from tensorflow.python.keras.saving.utils_v1 import signature_def_utils as unexported_signature_utils
|
|
from tensorflow.python.saved_model import signature_def_utils
|
|
|
|
|
|
class ExportOutput(object):
|
|
"""Represents an output of a model that can be served.
|
|
|
|
These typically correspond to model heads.
|
|
"""
|
|
|
|
__metaclass__ = abc.ABCMeta
|
|
|
|
_SEPARATOR_CHAR = '/'
|
|
|
|
@abc.abstractmethod
|
|
def as_signature_def(self, receiver_tensors):
|
|
"""Generate a SignatureDef proto for inclusion in a MetaGraphDef.
|
|
|
|
The SignatureDef will specify outputs as described in this ExportOutput,
|
|
and will use the provided receiver_tensors as inputs.
|
|
|
|
Args:
|
|
receiver_tensors: a `Tensor`, or a dict of string to `Tensor`, specifying
|
|
input nodes that will be fed.
|
|
"""
|
|
pass
|
|
|
|
def _check_output_key(self, key, error_label):
|
|
# For multi-head models, the key can be a tuple.
|
|
if isinstance(key, tuple):
|
|
key = self._SEPARATOR_CHAR.join(key)
|
|
|
|
if not isinstance(key, str):
|
|
raise ValueError(
|
|
'{} output key must be a string; got {}.'.format(error_label, key))
|
|
return key
|
|
|
|
def _wrap_and_check_outputs(
|
|
self, outputs, single_output_default_name, error_label=None):
|
|
"""Wraps raw tensors as dicts and checks type.
|
|
|
|
Note that we create a new dict here so that we can overwrite the keys
|
|
if necessary.
|
|
|
|
Args:
|
|
outputs: A `Tensor` or a dict of string to `Tensor`.
|
|
single_output_default_name: A string key for use in the output dict
|
|
if the provided `outputs` is a raw tensor.
|
|
error_label: descriptive string for use in error messages. If none,
|
|
single_output_default_name will be used.
|
|
|
|
Returns:
|
|
A dict of tensors
|
|
|
|
Raises:
|
|
ValueError: if the outputs dict keys are not strings or tuples of strings
|
|
or the values are not Tensors.
|
|
"""
|
|
if not isinstance(outputs, dict):
|
|
outputs = {single_output_default_name: outputs}
|
|
|
|
output_dict = {}
|
|
for key, value in outputs.items():
|
|
error_name = error_label or single_output_default_name
|
|
key = self._check_output_key(key, error_name)
|
|
if not isinstance(value, ops.Tensor):
|
|
raise ValueError(
|
|
'{} output value must be a Tensor; got {}.'.format(
|
|
error_name, value))
|
|
|
|
output_dict[key] = value
|
|
return output_dict
|
|
|
|
|
|
class ClassificationOutput(ExportOutput):
|
|
"""Represents the output of a classification head.
|
|
|
|
Either classes or scores or both must be set.
|
|
|
|
The classes `Tensor` must provide string labels, not integer class IDs.
|
|
|
|
If only classes is set, it is interpreted as providing top-k results in
|
|
descending order.
|
|
|
|
If only scores is set, it is interpreted as providing a score for every class
|
|
in order of class ID.
|
|
|
|
If both classes and scores are set, they are interpreted as zipped, so each
|
|
score corresponds to the class at the same index. Clients should not depend
|
|
on the order of the entries.
|
|
"""
|
|
|
|
def __init__(self, scores=None, classes=None):
|
|
"""Constructor for `ClassificationOutput`.
|
|
|
|
Args:
|
|
scores: A float `Tensor` giving scores (sometimes but not always
|
|
interpretable as probabilities) for each class. May be `None`, but
|
|
only if `classes` is set. Interpretation varies-- see class doc.
|
|
classes: A string `Tensor` giving predicted class labels. May be `None`,
|
|
but only if `scores` is set. Interpretation varies-- see class doc.
|
|
|
|
Raises:
|
|
ValueError: if neither classes nor scores is set, or one of them is not a
|
|
`Tensor` with the correct dtype.
|
|
"""
|
|
if (scores is not None
|
|
and not (isinstance(scores, ops.Tensor)
|
|
and scores.dtype.is_floating)):
|
|
raise ValueError('Classification scores must be a float32 Tensor; '
|
|
'got {}'.format(scores))
|
|
if (classes is not None
|
|
and not (isinstance(classes, ops.Tensor)
|
|
and dtypes.as_dtype(classes.dtype) == dtypes.string)):
|
|
raise ValueError('Classification classes must be a string Tensor; '
|
|
'got {}'.format(classes))
|
|
if scores is None and classes is None:
|
|
raise ValueError('At least one of scores and classes must be set.')
|
|
|
|
self._scores = scores
|
|
self._classes = classes
|
|
|
|
@property
|
|
def scores(self):
|
|
return self._scores
|
|
|
|
@property
|
|
def classes(self):
|
|
return self._classes
|
|
|
|
def as_signature_def(self, receiver_tensors):
|
|
if len(receiver_tensors) != 1:
|
|
raise ValueError('Classification input must be a single string Tensor; '
|
|
'got {}'.format(receiver_tensors))
|
|
(_, examples), = receiver_tensors.items()
|
|
if dtypes.as_dtype(examples.dtype) != dtypes.string:
|
|
raise ValueError('Classification input must be a single string Tensor; '
|
|
'got {}'.format(receiver_tensors))
|
|
return signature_def_utils.classification_signature_def(
|
|
examples, self.classes, self.scores)
|
|
|
|
|
|
class RegressionOutput(ExportOutput):
|
|
"""Represents the output of a regression head."""
|
|
|
|
def __init__(self, value):
|
|
"""Constructor for `RegressionOutput`.
|
|
|
|
Args:
|
|
value: a float `Tensor` giving the predicted values. Required.
|
|
|
|
Raises:
|
|
ValueError: if the value is not a `Tensor` with dtype tf.float32.
|
|
"""
|
|
if not (isinstance(value, ops.Tensor) and value.dtype.is_floating):
|
|
raise ValueError('Regression output value must be a float32 Tensor; '
|
|
'got {}'.format(value))
|
|
self._value = value
|
|
|
|
@property
|
|
def value(self):
|
|
return self._value
|
|
|
|
def as_signature_def(self, receiver_tensors):
|
|
if len(receiver_tensors) != 1:
|
|
raise ValueError('Regression input must be a single string Tensor; '
|
|
'got {}'.format(receiver_tensors))
|
|
(_, examples), = receiver_tensors.items()
|
|
if dtypes.as_dtype(examples.dtype) != dtypes.string:
|
|
raise ValueError('Regression input must be a single string Tensor; '
|
|
'got {}'.format(receiver_tensors))
|
|
return signature_def_utils.regression_signature_def(examples, self.value)
|
|
|
|
|
|
class PredictOutput(ExportOutput):
|
|
"""Represents the output of a generic prediction head.
|
|
|
|
A generic prediction need not be either a classification or a regression.
|
|
|
|
Named outputs must be provided as a dict from string to `Tensor`,
|
|
"""
|
|
_SINGLE_OUTPUT_DEFAULT_NAME = 'output'
|
|
|
|
def __init__(self, outputs):
|
|
"""Constructor for PredictOutput.
|
|
|
|
Args:
|
|
outputs: A `Tensor` or a dict of string to `Tensor` representing the
|
|
predictions.
|
|
|
|
Raises:
|
|
ValueError: if the outputs is not dict, or any of its keys are not
|
|
strings, or any of its values are not `Tensor`s.
|
|
"""
|
|
|
|
self._outputs = self._wrap_and_check_outputs(
|
|
outputs, self._SINGLE_OUTPUT_DEFAULT_NAME, error_label='Prediction')
|
|
|
|
@property
|
|
def outputs(self):
|
|
return self._outputs
|
|
|
|
def as_signature_def(self, receiver_tensors):
|
|
return signature_def_utils.predict_signature_def(receiver_tensors,
|
|
self.outputs)
|
|
|
|
|
|
class _SupervisedOutput(ExportOutput):
|
|
"""Represents the output of a supervised training or eval process."""
|
|
__metaclass__ = abc.ABCMeta
|
|
|
|
LOSS_NAME = 'loss'
|
|
PREDICTIONS_NAME = 'predictions'
|
|
METRICS_NAME = 'metrics'
|
|
|
|
METRIC_VALUE_SUFFIX = 'value'
|
|
METRIC_UPDATE_SUFFIX = 'update_op'
|
|
|
|
_loss = None
|
|
_predictions = None
|
|
_metrics = None
|
|
|
|
def __init__(self, loss=None, predictions=None, metrics=None):
|
|
"""Constructor for SupervisedOutput (ie, Train or Eval output).
|
|
|
|
Args:
|
|
loss: dict of Tensors or single Tensor representing calculated loss.
|
|
predictions: dict of Tensors or single Tensor representing model
|
|
predictions.
|
|
metrics: Dict of metric results keyed by name.
|
|
The values of the dict can be one of the following:
|
|
(1) instance of `Metric` class.
|
|
(2) (metric_value, update_op) tuples, or a single tuple.
|
|
metric_value must be a Tensor, and update_op must be a Tensor or Op.
|
|
|
|
Raises:
|
|
ValueError: if any of the outputs' dict keys are not strings or tuples of
|
|
strings or the values are not Tensors (or Operations in the case of
|
|
update_op).
|
|
"""
|
|
|
|
if loss is not None:
|
|
loss_dict = self._wrap_and_check_outputs(loss, self.LOSS_NAME)
|
|
self._loss = self._prefix_output_keys(loss_dict, self.LOSS_NAME)
|
|
if predictions is not None:
|
|
pred_dict = self._wrap_and_check_outputs(
|
|
predictions, self.PREDICTIONS_NAME)
|
|
self._predictions = self._prefix_output_keys(
|
|
pred_dict, self.PREDICTIONS_NAME)
|
|
if metrics is not None:
|
|
self._metrics = self._wrap_and_check_metrics(metrics)
|
|
|
|
def _prefix_output_keys(self, output_dict, output_name):
|
|
"""Prepend output_name to the output_dict keys if it doesn't exist.
|
|
|
|
This produces predictable prefixes for the pre-determined outputs
|
|
of SupervisedOutput.
|
|
|
|
Args:
|
|
output_dict: dict of string to Tensor, assumed valid.
|
|
output_name: prefix string to prepend to existing keys.
|
|
|
|
Returns:
|
|
dict with updated keys and existing values.
|
|
"""
|
|
|
|
new_outputs = {}
|
|
for key, val in output_dict.items():
|
|
key = self._prefix_key(key, output_name)
|
|
new_outputs[key] = val
|
|
return new_outputs
|
|
|
|
def _prefix_key(self, key, output_name):
|
|
if key.find(output_name) != 0:
|
|
key = output_name + self._SEPARATOR_CHAR + key
|
|
return key
|
|
|
|
def _wrap_and_check_metrics(self, metrics):
|
|
"""Handle the saving of metrics.
|
|
|
|
Metrics is either a tuple of (value, update_op), or a dict of such tuples.
|
|
Here, we separate out the tuples and create a dict with names to tensors.
|
|
|
|
Args:
|
|
metrics: Dict of metric results keyed by name.
|
|
The values of the dict can be one of the following:
|
|
(1) instance of `Metric` class.
|
|
(2) (metric_value, update_op) tuples, or a single tuple.
|
|
metric_value must be a Tensor, and update_op must be a Tensor or Op.
|
|
|
|
Returns:
|
|
dict of output_names to tensors
|
|
|
|
Raises:
|
|
ValueError: if the dict key is not a string, or the metric values or ops
|
|
are not tensors.
|
|
"""
|
|
if not isinstance(metrics, dict):
|
|
metrics = {self.METRICS_NAME: metrics}
|
|
|
|
outputs = {}
|
|
for key, value in metrics.items():
|
|
if isinstance(value, tuple):
|
|
metric_val, metric_op = value
|
|
else: # value is a keras.Metrics object
|
|
metric_val = value.result()
|
|
assert len(value.updates) == 1 # We expect only one update op.
|
|
metric_op = value.updates[0]
|
|
key = self._check_output_key(key, self.METRICS_NAME)
|
|
key = self._prefix_key(key, self.METRICS_NAME)
|
|
|
|
val_name = key + self._SEPARATOR_CHAR + self.METRIC_VALUE_SUFFIX
|
|
op_name = key + self._SEPARATOR_CHAR + self.METRIC_UPDATE_SUFFIX
|
|
if not isinstance(metric_val, ops.Tensor):
|
|
raise ValueError(
|
|
'{} output value must be a Tensor; got {}.'.format(
|
|
key, metric_val))
|
|
if not (tensor_util.is_tensor(metric_op) or
|
|
isinstance(metric_op, ops.Operation)):
|
|
raise ValueError(
|
|
'{} update_op must be a Tensor or Operation; got {}.'.format(
|
|
key, metric_op))
|
|
|
|
# We must wrap any ops (or variables) in a Tensor before export, as the
|
|
# SignatureDef proto expects tensors only. See b/109740581
|
|
metric_op_tensor = metric_op
|
|
if not isinstance(metric_op, ops.Tensor):
|
|
with ops.control_dependencies([metric_op]):
|
|
metric_op_tensor = constant_op.constant([], name='metric_op_wrapper')
|
|
|
|
outputs[val_name] = metric_val
|
|
outputs[op_name] = metric_op_tensor
|
|
|
|
return outputs
|
|
|
|
@property
|
|
def loss(self):
|
|
return self._loss
|
|
|
|
@property
|
|
def predictions(self):
|
|
return self._predictions
|
|
|
|
@property
|
|
def metrics(self):
|
|
return self._metrics
|
|
|
|
@abc.abstractmethod
|
|
def _get_signature_def_fn(self):
|
|
"""Returns a function that produces a SignatureDef given desired outputs."""
|
|
pass
|
|
|
|
def as_signature_def(self, receiver_tensors):
|
|
signature_def_fn = self._get_signature_def_fn()
|
|
return signature_def_fn(
|
|
receiver_tensors, self.loss, self.predictions, self.metrics)
|
|
|
|
|
|
class TrainOutput(_SupervisedOutput):
|
|
"""Represents the output of a supervised training process.
|
|
|
|
This class generates the appropriate signature def for exporting
|
|
training output by type-checking and wrapping loss, predictions, and metrics
|
|
values.
|
|
"""
|
|
|
|
def _get_signature_def_fn(self):
|
|
return unexported_signature_utils.supervised_train_signature_def
|
|
|
|
|
|
class EvalOutput(_SupervisedOutput):
|
|
"""Represents the output of a supervised eval process.
|
|
|
|
This class generates the appropriate signature def for exporting
|
|
eval output by type-checking and wrapping loss, predictions, and metrics
|
|
values.
|
|
"""
|
|
|
|
def _get_signature_def_fn(self):
|
|
return unexported_signature_utils.supervised_eval_signature_def
|
|
# LINT.ThenChange(//tensorflow/python/saved_model/model_utils/export_output.py)
|