Allow one to forward features to predictions
jankrynauw opened this issue · 3 comments
jankrynauw commented
We would like to forward a particular 'key' column which is part of the features to appear alongside the predictions - this is to be able to identify to which set of features a particular prediction belongs to. Here is an example of predictions output using the tensorflow.contrib.estimator.multi_class_head:
{"classes": ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"],
"scores": [0.06819603592157364, 0.0864366963505745, 0.12838752567768097, 0.046013250946998596, 0.03129083290696144, 0.1518409103155136, 0.1248951405286789, 0.15043732523918152, 0.0821763351559639, 0.13032598793506622]}
We would therefore like to add a key attribute to this prediction.
estimator = tf.contrib.estimator.forward_features(estimator, ['key'])
gives the following error:
The adanet.Estimator's model_fn should not be called directly in TRAIN mode, because its behavior is undefined outside the context of its train method.
The current workaround is to subclass the head
cweill commented
@jankrynauw Thank you for the feature request. This looks like it could be implemented in a way similar to adanet.Estimator(metric_fn=...) which we added in 53f5f5b.
Unfortunately, we don't have much time to add this ourselves, but if anyone has cycles to a PR, they are very welcome.
chamorajg commented
@jankrynauw Can you send me the code you were working on so that I can test it and work on it.
jankrynauw commented
See the predict section of the _create_tpu_estimator_spec method, as well as the predict_feature_keys=None parameter in the __init__ section:
# Implement a custom Head to control the prediction output. Most of this code is directly from
# `_MultiClassHeadWithSoftmaxCrossEntropyLoss` in `tensorflow_estimator.python.estimator.canned.head`
class MultiClassHead(_MultiClassHeadWithSoftmaxCrossEntropyLoss):
def __init__(self,
n_classes,
predict_feature_keys=None,
weight_column=None,
label_vocabulary=None,
loss_reduction=losses.Reduction.SUM_OVER_BATCH_SIZE,
loss_fn=None,
name=None):
"""Creates a '_Head' for multi class classification.
The head expects `logits` with shape `[D0, D1, ... DN, n_classes]`. In many applications, the shape is
`[batch_size, n_classes]`.
`labels` must be a dense `Tensor` with shape matching `logits`, namely
`[D0, D1, ... DN, 1]`. If `label_vocabulary` given, `labels` must be a string `Tensor` with values from the
vocabulary. If `label_vocabulary` is not given, `labels` must be an integer `Tensor` with values specifying the
class index.
If `weight_column` is specified, weights must be of shape `[D0, D1, ... DN]`, or `[D0, D1, ... DN, 1]`.
The loss is the weighted sum over the input dimensions. Namely, if the input labels have shape
`[batch_size, 1]`, the loss is the weighted sum over `batch_size`.
Also supports custom `loss_fn`. `loss_fn` takes `(labels, logits)` or `(labels, logits, features)` as arguments
and returns unreduced loss with shape `[D0, D1, ... DN, 1]`. `loss_fn` must support integer `labels` with shape
`[D0, D1, ... DN, 1]`. Namely, the head applies `label_vocabulary` to the input labels before passing them to
`loss_fn`.
Args:
n_classes: Number of classes, must be greater than 2 (for 2 classes,
use `_BinaryLogisticHeadWithSigmoidCrossEntropyLoss`).
predict_feature_keys: A `string` or a `list` of `string`. If it is `None`, all of the `features` in `dict`
is forwarded to the `predictions`. If it is a `string`, only given key is forwarded. If it is a `list`
of strings, all the given `keys` are forwarded.
weight_column: A string or a `_NumericColumn` created by `tf.feature_column.numeric_column` defining feature
column representing weights. It is used to down weight or boost examples during training. It will be
multiplied by the loss of the example.
label_vocabulary: A list or tuple of strings representing possible label values. If it is not given, that
means labels are already encoded as an integer within [0, n_classes). If given, labels must be of string
type and have any value in `label_vocabulary`. Note that errors will be raised if `label_vocabulary` is
not provided but labels are strings.
loss_reduction: One of `tf.losses.Reduction` except `NONE`. Describes how to reduce training loss over
batch. Defaults to `SUM_OVER_BATCH_SIZE`.
loss_fn: Optional loss function.
name: name of the head. If provided, summary and metrics keys will be suffixed by `"/" + name`. Also used as
`name_scope` when creating ops.
Returns:
An instance of `_Head` for multi class classification.
"""
self._predict_feature_keys = predict_feature_keys # Customisation done by alis.
super(MultiClassHead, self).__init__(n_classes=n_classes,
weight_column=weight_column,
label_vocabulary=label_vocabulary,
loss_reduction=loss_reduction,
loss_fn=loss_fn,
name=name)
def _create_tpu_estimator_spec(self,
features,
mode,
logits,
labels=None,
optimizer=None,
train_op_fn=None,
regularization_losses=None):
"""Returns a `model_fn._TPUEstimatorSpec`.
Args:
features: Input `dict` of `Tensor` or `SparseTensor` objects.
mode: Estimator's `ModeKeys`.
logits: logits `Tensor` with shape `[D0, D1, ... DN, logits_dimension]`. For many applications, the shape is
`[batch_size, logits_dimension]`.
labels: Labels integer or string `Tensor` with shape matching `logits`, namely `[D0, D1, ... DN, 1]` or
`[D0, D1, ... DN]`. `labels` is required argument when `mode` equals `TRAIN` or `EVAL`.
optimizer: `Optimizer` instance to optimize the loss in TRAIN mode. Namely, sets `train_op =
optimizer.minimize(loss, global_step)`, which updates variables and increments `global_step`.
train_op_fn: Function that takes a scalar loss `Tensor` and returns `train_op`. Used if `optimizer` is `None`.
regularization_losses: A list of additional scalar losses to be added to the training loss, such as
regularization losses. These losses are usually expressed as a batch average, so for best results users
need to set `loss_reduction=SUM_OVER_BATCH_SIZE` or `loss_reduction=SUM_OVER_NONZERO_WEIGHTS` when
creating the head to avoid scaling errors.
Returns:
A `model_fn._TPUEstimatorSpec` instance.
"""
with ops.name_scope(self._name, 'head'):
logits = _check_logits_final_dim(logits, self.logits_dimension)
with ops.name_scope(None, 'predictions', (logits,)):
# Get the predicted class IDs and names
class_ids = math_ops.argmax(logits, axis=-1, name=PredictionKeys.CLASS_IDS)
class_ids = array_ops.expand_dims(class_ids, axis=-1)
if self._label_vocabulary:
table = lookup_ops.index_to_string_table_from_tensor(vocabulary_list=self._label_vocabulary,
name='class_string_lookup')
classes = table.lookup(class_ids)
else:
classes = string_ops.as_string(class_ids, name='str_classes')
# Compute the predicted probabilities
probabilities = nn.softmax(logits, name=PredictionKeys.PROBABILITIES)
# Compute the predicted score we use for ranking
predicted_class = math_ops.argmax(logits, axis=-1) + 1
predicted_class_probability = math_ops.reduce_max(probabilities, axis=-1)
predicted = tf.expand_dims(tf.cast(predicted_class, tf.float32) - predicted_class_probability, axis=-1)
predictions = {
PredictionKeys.CLASS_IDS: class_ids,
PredictionKeys.CLASSES: classes,
PredictionKeys.LOGITS: logits,
PredictionKeys.PREDICTED: predicted,
PredictionKeys.PROBABILITIES: probabilities,
}
# Predict
if mode == model_fn.ModeKeys.PREDICT:
# CUSTOMISATION DONE HERE. Attach the feature keys to the predictions
for key in self._predict_feature_keys:
predictions[key] = features[key]
classifier_output = _classification_output(scores=probabilities,
n_classes=self._n_classes,
label_vocabulary=self._label_vocabulary)
return model_fn._TPUEstimatorSpec(mode=model_fn.ModeKeys.PREDICT,
predictions=predictions,
export_outputs={
_DEFAULT_SERVING_KEY: export_output.PredictOutput(predictions),
_CLASSIFY_SERVING_KEY: classifier_output,
_PREDICT_SERVING_KEY: export_output.PredictOutput(predictions)
})
# Compute loss
training_loss, unreduced_loss, weights, label_ids = self.create_loss(features=features,
mode=mode,
logits=logits,
labels=labels)
if regularization_losses:
regularization_loss = math_ops.add_n(regularization_losses)
regularized_training_loss = math_ops.add_n([training_loss, regularization_loss])
else:
regularization_loss = None
regularized_training_loss = training_loss
# Eval
if mode == model_fn.ModeKeys.EVAL:
return model_fn._TPUEstimatorSpec(mode=model_fn.ModeKeys.EVAL,
predictions=predictions,
loss=regularized_training_loss,
eval_metrics=_create_eval_metrics_tuple(self._eval_metric_ops, {
'labels': label_ids,
'class_ids': class_ids,
'weights': weights,
'unreduced_loss': unreduced_loss,
'regularization_loss': regularization_loss
}))
# Train
if optimizer is not None:
if train_op_fn is not None:
raise ValueError('train_op_fn and optimizer cannot both be set.')
train_op = optimizer.minimize(regularized_training_loss, global_step=training_util.get_global_step())
elif train_op_fn is not None:
train_op = train_op_fn(regularized_training_loss)
else:
raise ValueError('train_op_fn and optimizer cannot both be None.')
train_op = _append_update_ops(train_op)
# Only summarize mean_loss for SUM reduction to preserve backwards compatibility. Otherwise skip it to
# avoid unnecessary computation.
if self._loss_reduction == losses.Reduction.SUM:
example_weight_sum = math_ops.reduce_sum(weights * array_ops.ones_like(unreduced_loss))
mean_loss = training_loss / example_weight_sum
else:
mean_loss = None
with ops.name_scope(''):
keys = metric_keys.MetricKeys
summary.scalar(_summary_key(self._name, keys.LOSS), regularized_training_loss)
if mean_loss is not None:
summary.scalar(_summary_key(self._name, keys.LOSS_MEAN), mean_loss)
if regularization_loss is not None:
summary.scalar(_summary_key(self._name, keys.LOSS_REGULARIZATION), regularization_loss)
return model_fn._TPUEstimatorSpec(mode=model_fn.ModeKeys.TRAIN,
predictions=predictions,
loss=regularized_training_loss,
train_op=train_op)