# Copyright 2018 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. # ============================================================================== """Regression head.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function import tensorflow as tf from tensorflow.python.framework import ops from tensorflow_estimator.python.estimator import model_fn from tensorflow_estimator.python.estimator.canned import metric_keys from tensorflow_estimator.python.estimator.canned import prediction_keys from tensorflow_estimator.python.estimator.estimator_export import estimator_export from tensorflow_estimator.python.estimator.export import export_output from tensorflow_estimator.python.estimator.head import base_head from tensorflow_estimator.python.estimator.mode_keys import ModeKeys @estimator_export('estimator.RegressionHead') class RegressionHead(base_head.Head): """Creates a `Head` for regression using the `mean_squared_error` loss. The loss is the weighted sum over all input dimensions. Namely, if the input labels have shape `[batch_size, label_dimension]`, the loss is the weighted sum over both `batch_size` and `label_dimension`. The head expects `logits` with shape `[D0, D1, ... DN, label_dimension]`. In many applications, the shape is `[batch_size, label_dimension]`. The `labels` shape must match `logits`, namely `[D0, D1, ... DN, label_dimension]`. If `label_dimension=1`, shape `[D0, D1, ... DN]` is also supported. If `weight_column` is specified, weights must be of shape `[D0, D1, ... DN]`, `[D0, D1, ... DN, 1]` or `[D0, D1, ... DN, label_dimension]`. Supports custom `loss_fn`. `loss_fn` takes `(labels, logits)` or `(labels, logits, features, loss_reduction)` as arguments and returns unreduced loss with shape `[D0, D1, ... DN, label_dimension]`. Also supports custom `inverse_link_fn`, also known as 'mean function'. `inverse_link_fn` is only used in `PREDICT` mode. It takes `logits` as argument and returns predicted values. This function is the inverse of the link function defined in https://en.wikipedia.org/wiki/Generalized_linear_model#Link_function Namely, for poisson regression, set `inverse_link_fn=tf.exp`. Usage: >>> head = tf.estimator.RegressionHead() >>> logits = np.array(((45,), (41,),), dtype=np.float32) >>> labels = np.array(((43,), (44,),), dtype=np.int32) >>> features = {'x': np.array(((42,),), dtype=np.float32)} >>> # expected_loss = weighted_loss / batch_size >>> # = (43-45)^2 + (44-41)^2 / 2 = 6.50 >>> loss = head.loss(labels, logits, features=features) >>> print('{:.2f}'.format(loss.numpy())) 6.50 >>> eval_metrics = head.metrics() >>> updated_metrics = head.update_metrics( ... eval_metrics, features, logits, labels) >>> for k in sorted(updated_metrics): ... print('{} : {:.2f}'.format(k, updated_metrics[k].result().numpy())) average_loss : 6.50 label/mean : 43.50 prediction/mean : 43.00 >>> preds = head.predictions(logits) >>> print(preds['predictions']) tf.Tensor( [[45.] [41.]], shape=(2, 1), dtype=float32) Usage with a canned estimator: ```python my_head = tf.estimator.RegressionHead() my_estimator = tf.estimator.DNNEstimator( head=my_head, hidden_units=..., feature_columns=...) ``` It can also be used with a custom `model_fn`. Example: ```python def _my_model_fn(features, labels, mode): my_head = tf.estimator.RegressionHead() logits = tf.keras.Model(...)(features) return my_head.create_estimator_spec( features=features, mode=mode, labels=labels, optimizer=tf.keras.optimizers.Adagrad(lr=0.1), logits=logits) my_estimator = tf.estimator.Estimator(model_fn=_my_model_fn) ``` Args: 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_dimension: Number of regression labels per example. This is the size of the last dimension of the labels `Tensor` (typically, this has shape `[batch_size, label_dimension]`). loss_reduction: One of `tf.losses.Reduction` except `NONE`. Decides how to reduce training loss over batch and label dimension. Defaults to `SUM_OVER_BATCH_SIZE`, namely weighted sum of losses divided by `batch_size * label_dimension`. loss_fn: Optional loss function. Defaults to `mean_squared_error`. inverse_link_fn: Optional inverse link function, also known as 'mean function'. Defaults to identity. name: name of the head. If provided, summary and metrics keys will be suffixed by `"/" + name`. Also used as `name_scope` when creating ops. """ def __init__(self, label_dimension=1, weight_column=None, loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE, loss_fn=None, inverse_link_fn=None, name=None): if label_dimension < 1: raise ValueError('Invalid label_dimension {}.'.format(label_dimension)) base_head.validate_loss_reduction(loss_reduction) if loss_fn: base_head.validate_loss_fn_args(loss_fn) self._logits_dimension = label_dimension self._weight_column = weight_column self._loss_reduction = loss_reduction self._loss_fn = loss_fn self._inverse_link_fn = inverse_link_fn self._name = name # Metric keys. keys = metric_keys.MetricKeys self._loss_mean_key = self._summary_key(keys.LOSS_MEAN) self._prediction_mean_key = self._summary_key(keys.PREDICTION_MEAN) self._label_mean_key = self._summary_key(keys.LABEL_MEAN) self._loss_regularization_key = self._summary_key(keys.LOSS_REGULARIZATION) @property def name(self): """See `base_head.Head` for details.""" return self._name @property def logits_dimension(self): """See `base_head.Head` for details.""" return self._logits_dimension @property def loss_reduction(self): """See `base_head.Head` for details.""" return self._loss_reduction def _processed_labels(self, logits, labels): labels = base_head.check_dense_labels_match_logits_and_reshape( labels=labels, logits=logits, expected_labels_dimension=self._logits_dimension) labels = tf.cast(labels, dtype=tf.dtypes.float32) return labels def _unweighted_loss_and_weights(self, logits, labels, features): """Computes unweighted loss and weights.""" if self._loss_fn: unweighted_loss = base_head.call_loss_fn( loss_fn=self._loss_fn, labels=labels, logits=logits, features=features, expected_loss_dim=self._logits_dimension) else: unweighted_loss = tf.compat.v1.losses.mean_squared_error( labels=labels, predictions=logits, reduction=tf.compat.v1.losses.Reduction.NONE) weights = base_head.get_weights_and_check_match_logits( features=features, weight_column=self._weight_column, logits=logits, allow_per_logit_weights=True) return unweighted_loss, weights def loss(self, labels, logits, features=None, mode=None, regularization_losses=None): """Return predictions based on keys. See `base_head.Head` for details.""" del mode # Unused for this head. with ops.name_scope( 'losses', values=(logits, labels, regularization_losses, features)): logits = base_head.check_logits_final_dim(logits, self._logits_dimension) labels = self._processed_labels(logits, labels) unweighted_loss, weights = self._unweighted_loss_and_weights( logits, labels, features) training_loss = tf.compat.v2.keras.__internal__.losses.compute_weighted_loss( unweighted_loss, sample_weight=weights, reduction=self._loss_reduction) regularization_loss = tf.math.add_n( regularization_losses) if regularization_losses is not None else None regularized_training_loss = ( training_loss + regularization_loss if regularization_loss is not None else training_loss) return regularized_training_loss def predictions(self, logits): """Return predictions based on keys. See `base_head.Head` for details. Args: logits: logits `Tensor` with shape `[D0, D1, ... DN, logits_dimension]`. For many applications, the shape is `[batch_size, logits_dimension]`. Returns: A dict of predictions. """ logits = base_head.check_logits_final_dim(logits, self._logits_dimension) pred_keys = prediction_keys.PredictionKeys with ops.name_scope('predictions', values=(logits,)): if self._inverse_link_fn: predicted_value = self._inverse_link_fn(logits) predictions = { pred_keys.PREDICTIONS: predicted_value, pred_keys.LOGITS: logits, } else: predicted_value = logits predictions = {pred_keys.PREDICTIONS: predicted_value} return predictions def metrics(self, regularization_losses=None): """Creates metrics. See `base_head.Head` for details.""" with ops.name_scope('metrics', values=(regularization_losses,)): keys = metric_keys.MetricKeys eval_metrics = {} eval_metrics[self._loss_mean_key] = tf.keras.metrics.Mean( name=keys.LOSS_MEAN) eval_metrics[self._prediction_mean_key] = tf.keras.metrics.Mean( name=keys.PREDICTION_MEAN) eval_metrics[self._label_mean_key] = tf.keras.metrics.Mean( name=keys.LABEL_MEAN) if regularization_losses is not None: eval_metrics[self._loss_regularization_key] = tf.keras.metrics.Mean( name=keys.LOSS_REGULARIZATION) return eval_metrics def update_metrics(self, eval_metrics, features, logits, labels, regularization_losses=None): """Updates eval metrics. See `base_head.Head` for details.""" # Compute predictions. predictions = self.predictions(logits) predicted_value = predictions[prediction_keys.PredictionKeys.PREDICTIONS] logits = base_head.check_logits_final_dim(logits, self.logits_dimension) label_ids = self._processed_labels(logits, labels) unweighted_loss, weights = self._unweighted_loss_and_weights( logits, label_ids, features) # Update metrics. eval_metrics[self._loss_mean_key].update_state( values=unweighted_loss, sample_weight=weights) eval_metrics[self._label_mean_key].update_state( values=labels, sample_weight=weights) base_head.update_metric_with_broadcast_weights( eval_metrics[self._prediction_mean_key], predicted_value, weights) if regularization_losses is not None: regularization_loss = tf.math.add_n(regularization_losses) eval_metrics[self._loss_regularization_key].update_state( values=regularization_loss) return eval_metrics def _create_tpu_estimator_spec(self, features, mode, logits, labels=None, optimizer=None, trainable_variables=None, train_op_fn=None, update_ops=None, regularization_losses=None): """Returns an `EstimatorSpec`. Args: features: Input `dict` mapping string feature names to `Tensor` or `SparseTensor` objects containing the values for that feature in a minibatch. Often to be used to fetch example-weight tensor. 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 `Tensor` with shape matching `logits`, namely `[D0, D1, ... DN, logits_dimension]`. When `logits_dimension=1`, shape `[D0, D1, ... DN]` is also supported. `labels` is a required argument when `mode` equals `TRAIN` or `EVAL`. optimizer: An `tf.keras.optimizers.Optimizer` instance to optimize the loss in TRAIN mode. Namely, sets `train_op = optimizer.get_updates(loss, trainable_variables)`, which updates variables to minimize `loss`. trainable_variables: A list or tuple of `Variable` objects to update to minimize `loss`. In Tensorflow 1.x, by default these are the list of variables collected in the graph under the key `GraphKeys.TRAINABLE_VARIABLES`. As Tensorflow 2.x doesn't have collections and GraphKeys, trainable_variables need to be passed explicitly here. train_op_fn: Function that takes a scalar loss `Tensor` and returns `train_op`. Used if `optimizer` is `None`. update_ops: A list or tuple of update ops to be run at training time. For example, layers such as BatchNormalization create mean and variance update ops that need to be run at training time. In Tensorflow 1.x, these are thrown into an UPDATE_OPS collection. As Tensorflow 2.x doesn't have collections, update_ops need to be passed explicitly here. 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` when creating the head to avoid scaling errors. Returns: A `model_fn._TPUEstimatorSpec` instance. Raises: ValueError: If both `train_op_fn` and `optimizer` are `None` in TRAIN mode, or if both are set. """ with ops.name_scope(self._name, 'head'): # Predict. predictions = self.predictions(logits) if mode == ModeKeys.PREDICT: keys = prediction_keys.PredictionKeys regression_output = export_output.RegressionOutput( value=predictions[keys.PREDICTIONS]) return model_fn._TPUEstimatorSpec( # pylint: disable=protected-access mode=ModeKeys.PREDICT, predictions=predictions, export_outputs={ base_head.DEFAULT_SERVING_KEY: regression_output, base_head.REGRESS_SERVING_KEY: regression_output, base_head.PREDICT_SERVING_KEY: export_output.PredictOutput( predictions) }) regularized_training_loss = self.loss( logits=logits, labels=labels, features=features, mode=mode, regularization_losses=regularization_losses) # Eval. if mode == ModeKeys.EVAL: eval_metrics = self.metrics(regularization_losses=regularization_losses) return model_fn._TPUEstimatorSpec( # pylint: disable=protected-access mode=ModeKeys.EVAL, predictions=predictions, loss=regularized_training_loss, eval_metrics=base_head.create_eval_metrics_tuple( self.update_metrics, { 'eval_metrics': eval_metrics, 'features': features, 'logits': logits, 'labels': labels, 'regularization_losses': regularization_losses })) # Train. train_op = base_head.create_estimator_spec_train_op( head_name=self._name, optimizer=optimizer, train_op_fn=train_op_fn, update_ops=update_ops, trainable_variables=trainable_variables, regularized_training_loss=regularized_training_loss, loss_reduction=self._loss_reduction) # Create summary. base_head.create_estimator_spec_summary( regularized_training_loss=regularized_training_loss, regularization_losses=regularization_losses, summary_key_fn=self._summary_key) return model_fn._TPUEstimatorSpec( # pylint: disable=protected-access mode=ModeKeys.TRAIN, predictions=predictions, loss=regularized_training_loss, train_op=train_op) @estimator_export('estimator.PoissonRegressionHead') class PoissonRegressionHead(RegressionHead): """Creates a `Head` for poisson regression using `tf.nn.log_poisson_loss`. The loss is the weighted sum over all input dimensions. Namely, if the input labels have shape `[batch_size, label_dimension]`, the loss is the weighted sum over both `batch_size` and `label_dimension`. The head expects `logits` with shape `[D0, D1, ... DN, label_dimension]`. In many applications, the shape is `[batch_size, label_dimension]`. The `labels` shape must match `logits`, namely `[D0, D1, ... DN, label_dimension]`. If `label_dimension=1`, shape `[D0, D1, ... DN]` is also supported. If `weight_column` is specified, weights must be of shape `[D0, D1, ... DN]`, `[D0, D1, ... DN, 1]` or `[D0, D1, ... DN, label_dimension]`. This is implemented as a generalized linear model, see https://en.wikipedia.org/wiki/Generalized_linear_model. The head can be used with a canned estimator. Example: ```python my_head = tf.estimator.PoissonRegressionHead() my_estimator = tf.estimator.DNNEstimator( head=my_head, hidden_units=..., feature_columns=...) ``` It can also be used with a custom `model_fn`. Example: ```python def _my_model_fn(features, labels, mode): my_head = tf.estimator.PoissonRegressionHead() logits = tf.keras.Model(...)(features) return my_head.create_estimator_spec( features=features, mode=mode, labels=labels, optimizer=tf.keras.optimizers.Adagrad(lr=0.1), logits=logits) my_estimator = tf.estimator.Estimator(model_fn=_my_model_fn) ``` Args: 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_dimension: Number of regression labels per example. This is the size of the last dimension of the labels `Tensor` (typically, this has shape `[batch_size, label_dimension]`). loss_reduction: One of `tf.losses.Reduction` except `NONE`. Decides how to reduce training loss over batch and label dimension. Defaults to `SUM_OVER_BATCH_SIZE`, namely weighted sum of losses divided by `batch size * label_dimension`. compute_full_loss: Whether to include the constant `log(z!)` term in computing the poisson loss. See `tf.nn.log_poisson_loss` for the full documentation. name: name of the head. If provided, summary and metrics keys will be suffixed by `"/" + name`. Also used as `name_scope` when creating ops. """ def __init__(self, label_dimension=1, weight_column=None, loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE, compute_full_loss=True, name=None): self._compute_full_loss = compute_full_loss super(PoissonRegressionHead, self).__init__( label_dimension=label_dimension, weight_column=weight_column, loss_reduction=loss_reduction, loss_fn=self._poisson_loss, inverse_link_fn=tf.math.exp, name=name) def _poisson_loss(self, labels, logits): return tf.nn.log_poisson_loss( targets=labels, log_input=logits, compute_full_loss=self._compute_full_loss) @estimator_export('estimator.LogisticRegressionHead') class LogisticRegressionHead(RegressionHead): """Creates a `Head` for logistic regression. Uses `sigmoid_cross_entropy_with_logits` loss, which is the same as `BinaryClassHead`. The differences compared to `BinaryClassHead` are: * Does not support `label_vocabulary`. Instead, labels must be float in the range [0, 1]. * Does not calculate some metrics that do not make sense, such as AUC. * In `PREDICT` mode, only returns logits and predictions (`=tf.sigmoid(logits)`), whereas `BinaryClassHead` also returns probabilities, classes, and class_ids. * Export output defaults to `RegressionOutput`, whereas `BinaryClassHead` defaults to `PredictOutput`. The head expects `logits` with shape `[D0, D1, ... DN, 1]`. In many applications, the shape is `[batch_size, 1]`. The `labels` shape must match `logits`, namely `[D0, D1, ... DN]` or `[D0, D1, ... DN, 1]`. If `weight_column` is specified, weights must be of shape `[D0, D1, ... DN]` or `[D0, D1, ... DN, 1]`. This is implemented as a generalized linear model, see https://en.wikipedia.org/wiki/Generalized_linear_model. The head can be used with a canned estimator. Example: ```python my_head = tf.estimator.LogisticRegressionHead() my_estimator = tf.estimator.DNNEstimator( head=my_head, hidden_units=..., feature_columns=...) ``` It can also be used with a custom `model_fn`. Example: ```python def _my_model_fn(features, labels, mode): my_head = tf.estimator.LogisticRegressionHead() logits = tf.keras.Model(...)(features) return my_head.create_estimator_spec( features=features, mode=mode, labels=labels, optimizer=tf.keras.optimizers.Adagrad(lr=0.1), logits=logits) my_estimator = tf.estimator.Estimator(model_fn=_my_model_fn) ``` Args: 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. loss_reduction: One of `tf.losses.Reduction` except `NONE`. Decides how to reduce training loss over batch and label dimension. Defaults to `SUM_OVER_BATCH_SIZE`, namely weighted sum of losses divided by `batch size * label_dimension`. name: name of the head. If provided, summary and metrics keys will be suffixed by `"/" + name`. Also used as `name_scope` when creating ops. """ def _logistic_loss(self, labels, logits): labels = base_head.check_label_range( labels, n_classes=2, message='Labels must be in range [0, 1]') return tf.compat.v1.nn.sigmoid_cross_entropy_with_logits( labels=labels, logits=logits) def __init__(self, weight_column=None, loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE, name=None): super(LogisticRegressionHead, self).__init__( label_dimension=1, weight_column=weight_column, loss_reduction=loss_reduction, loss_fn=self._logistic_loss, inverse_link_fn=tf.math.sigmoid, name=name)