tf.contrib.tpu.TPUEstimatorSpec

Class TPUEstimatorSpec

Defined in tensorflow/contrib/tpu/python/tpu/tpu_estimator.py.

Ops and objects returned from a model_fn and passed to TPUEstimator.

See EstimatorSpec for mode, 'predictions, 'loss', 'train_op', and 'export_outputs`.

For evaluation, eval_metricsis a tuple of metric_fn and tensors, where metric_fn runs on CPU to generate metrics and tensors represents the Tensors transferred from TPU system to CPU host and passed to metric_fn. To be precise, TPU evaluation expects a slightly different signature from the tf.estimator.Estimator. While EstimatorSpec.eval_metric_ops expects a dict, TPUEstimatorSpec.eval_metrics is a tuple of metric_fn and tensors. The tensors could be a list of Tensors or dict of names to Tensors. The tensors usually specify the model logits, which are transferred back from TPU system to CPU host. All tensors must have be batch-major, i.e., the batch size is the first dimension. Once all tensors are available at CPU host from all shards, they are concatenated (on CPU) and passed as positional arguments to the metric_fn if tensors is list or keyword arguments if tensors is dict. metric_fn takes the tensors and returns a dict from metric string name to the result of calling a metric function, namely a (metric_tensor, update_op) tuple. See TPUEstimator for MNIST example how to specify the eval_metrics.

scaffold_fn is a function running on CPU to generate the Scaffold. This function should not capture any Tensors in model_fn.

host_call is a tuple of a function and a list or dictionary of tensors to pass to that function and returns a list of Tensors. host_call currently works for train() and evaluate(). The Tensors returned by the function is executed on the CPU on every step, so there is communication overhead when sending tensors from TPU to CPU. To reduce the overhead, try reducing the size of the tensors. The tensors are concatenated along their major (batch) dimension, and so must be >= rank 1. The host_call is useful for writing summaries with tf.contrib.summary.create_file_writer.

Properties

eval_metrics

Alias for field number 4

export_outputs

Alias for field number 5

host_call

Alias for field number 7

loss

Alias for field number 2

mode

Alias for field number 0

predictions

Alias for field number 1

scaffold_fn

Alias for field number 6

train_op

Alias for field number 3

Methods

__new__

@staticmethod
__new__(
    cls,
    mode,
    predictions=None,
    loss=None,
    train_op=None,
    eval_metrics=None,
    export_outputs=None,
    scaffold_fn=None,
    host_call=None
)

Creates a validated TPUEstimatorSpec instance.

as_estimator_spec

as_estimator_spec()

Creates an equivalent EstimatorSpec used by CPU train/eval.