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


Interface for the head/top of a model.

Head sits on top of the model network and handles computing the outputs of the network. Given logits (or output of a hidden layer), a Head knows how to compute predictions, loss, train_op, metrics and export outputs. It is meant to:

  1. Simplify writing model_fn and to make model_fn more configurable for Estimator.
  2. Simpilfy creating loss and metrics for the train and test loop in Eager execution.
  3. Support wide range of machine learning models. Since most heads can work with logits, they can support DNN, RNN, Wide, Wide&Deep, Global objectives, Gradient boosted trees and many other types of machine learning models.

Common usage:

Here is simplified model_fn to build a DNN regression model.

  def _my_dnn_model_fn(features, labels, mode, params, config=None):
    # Optionally your callers can pass head to model_fn as a param.
    head = tf.estimator.RegressionHead(...)

    feature_columns = tf.feature_column.numeric_column(...)
    feature_layer = tf.keras.layers.DenseFeatures(feature_columns)
    inputs = feature_layer(features)

    # Compute logits with tf.keras.layers API
    hidden_layer0 = tf.keras.layers.Dense(
        units=1000, activation="relu")(inputs)
    hidden_layer1 = tf.keras.layers.Dense(
        units=500, activation="relu")(hidden_layer0)
    logits = tf.keras.layers.Dense(
        units=head.logits_dimension, activation=None)(hidden_layer1)

    # Or use Keras model for logits computation
    model = tf.keras.Sequential()
    model.add(tf.keras.layers.Dense(units=1000, activation="relu"))
    model.add(tf.keras.layers.Dense(units=500, activation="relu"))
       units=head.logits_dimension, activation=None))
    logits = model(inputs)

    return head.create_estimator_spec(

logits_dimension Size of the last dimension of the logits Tensor.

Often is the number of classes, labels, or real values to be predicted. Typically, logits is of shape [batch_size, logits_dimension].

loss_reduction One of tf.losses.Reduction.

Describes how to reduce training loss over batch, such as mean or sum.

name The name of this head.



View source

Returns EstimatorSpec that a model_fn can return.

It is recommended to pass all args via name.

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 to be used by the head.
labels Labels Tensor, or dict mapping string label names to Tensor objects of the label values.
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 und