ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more


Creates a Head for regression using the mean_squared_error loss.

Inherits From: Head

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 Namely, for poisson regression, set inverse_link_fn=tf.exp.


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)
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)
   [41.]], shape=(2, 1), dtype=float32)

Usage with a canned estimator:

my_head = tf.estimator.RegressionHead()
my_estimator = tf.estimator.DNNEstimator(

It can also be used with a custom model_fn. Example:

def _my_model_fn(features, labels, mode):
  my_head = tf.estimator.RegressionHead()
  logits = tf.keras.Model(...)(features)

  return my_head.create_estimator_spec(

my_estimator = tf.estimator.Estimator(model_fn=_my_model_fn)

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 trainin