支持 TPU 的估算器。

警告：不建议将估算器用于新代码。估算器运行tf.compat.v1.Session-style 代码更难正确编写，并且可能出现意外行为，尤其是与 TF 2 代码结合使用时。估算器确实属于我们的兼容性保证，但不会收到除安全漏洞以外的任何修复。见迁移指南详情。

继承自：Estimator

用法

tf.compat.v1.estimator.tpu.TPUEstimator(
    model_fn=None, model_dir=None, config=None, params=None, use_tpu=True,
    train_batch_size=None, eval_batch_size=None, predict_batch_size=None,
    batch_axis=None, eval_on_tpu=True, export_to_tpu=True, export_to_cpu=True,
    warm_start_from=None, embedding_config_spec=None,
    export_saved_model_api_version=ExportSavedModelApiVersion.V1
)

迁移到 TF2

警告：这个 API 是为 TensorFlow v1 设计的。继续阅读有关如何从该 API 迁移到本机 TensorFlow v2 等效项的详细信息。见TensorFlow v1 到 TensorFlow v2 迁移指南有关如何迁移其余代码的说明。

TPU Estimator 管理自己的 TensorFlow 图和会话，因此与 TF2 行为不兼容。我们建议您迁移到较新的 tf.distribute.TPUStrategy 。有关详细信息，请参阅 TPU 指南。

TPUEstimator 还支持在 CPU 和 GPU 上进行训练。您无需定义单独的 tf.estimator.Estimator 。

TPUEstimator 处理在 TPU 设备上运行的许多细节，例如为每个内核复制输入和模型，以及定期返回主机以运行挂钩。

TPUEstimator 在调用 input_fn 和 model_fn 时将参数中的全局批量大小转换为 per-shard 批量大小。用户需要在构造函数中指定全局batch size，然后通过params['batch_size']获取input_fn和model_fn中每个shard的batch size。

• 对于训练，model_fn 得到 per-core 批量大小； input_fn 可能会得到 per-core 或 per-host 批量大小，具体取决于 TPUConfig 中的 per_host_input_for_training(有关详细信息，请参阅 TPUConfig 的文档字符串)。

• 对于评估和预测，model_fn 获取 per-core 批量大小，input_fn 获取 per-host 批量大小。

评估

model_fn 应该返回 TPUEstimatorSpec ，它需要 eval_metrics 用于 TPU 评估。如果eval_on_tpu 为 False，评估将在 CPU 或 GPU 上执行；在这种情况下，以下关于 TPU 评估的讨论不适用。

TPUEstimatorSpec.eval_metrics 是 metric_fn 和 tensors 的元组，其中 tensors 可以是 Tensor 的任何嵌套结构的列表(有关详细信息，请参见 TPUEstimatorSpec)。 metric_fn 接受 tensors 并将度量字符串名称中的 dict 返回到调用度量函数的结果，即 (metric_tensor, update_op) 元组。

可以将use_tpu设置为False进行测试。所有训练、评估和预测都将在 CPU 上执行。 input_fn 和 model_fn 将收到 train_batch_size 或 eval_batch_size 未修改为 params['batch_size'] 。

当前限制：

---

1. TPU 评估仅适用于单个主机(一个 TPU 工作器)，广播模式除外。

2. 用于评估的input_fn 不应引发end-of-input 异常(OutOfRangeError 或StopIteration)。并且所有评估步骤和所有批次都应该具有相同的大小。

示例(MNIST)：

# The metric Fn which runs on CPU.
def metric_fn(labels, logits):
  predictions = tf.argmax(logits, 1)
  return {
    'accuracy':tf.compat.v1.metrics.precision(
        labels=labels, predictions=predictions),
  }

# Your model Fn which runs on TPU (eval_metrics is list in this example)
def model_fn(features, labels, mode, config, params):
  ...
  logits = ...

  if mode = tf.estimator.ModeKeys.EVAL:
    return tpu_estimator.TPUEstimatorSpec(
        mode=mode,
        loss=loss,
        eval_metrics=(metric_fn, [labels, logits]))

# or specify the eval_metrics tensors as dict.
def model_fn(features, labels, mode, config, params):
  ...
  final_layer_output = ...

  if mode = tf.estimator.ModeKeys.EVAL:
    return tpu_estimator.TPUEstimatorSpec(
        mode=mode,
        loss=loss,
        eval_metrics=(metric_fn, {
            'labels':labels,
            'logits':final_layer_output,
        }))

预测

TPU 上的预测是支持大批量推理的实验性函数。它不是为latency-critical 系统设计的。此外，由于一些可用性问题，对于小数据集的预测，CPU .predict ，即使用 use_tpu=False 创建一个新的 TPUEstimator 实例可能更方便。

注意：与 TPU 训练/评估相比，input_fn用于预测应该引发 end-of-input 异常(OutOfRangeError或者StopIteration)，作为停止信号TPUEstimator.准确地说，由input_fn产生一批数据。这predict()API 一次处理一批。当到达数据源的末尾时，这些操作之一应引发end-of-input 异常。用户通常不需要手动执行此操作。只要数据集不会永远重复，tf.dataAPI 将在生产最后一批后自动引发 end-of-input 异常。

注意：Estimator.predict 返回一个 Python 生成器。请消耗来自生成器的所有数据，以便 TPUEstimator 可以为用户正确关闭 TPU 系统。

当前限制：

---

1. TPU 预测仅适用于单个主机(一个 TPU 工作器)。

2. input_fn 必须返回 Dataset 实例而不是 features 。事实上，.train() 和 .evaluate() 也支持 Dataset 作为返回值。

示例(MNIST)：

height = 32
width = 32
total_examples = 100

def predict_input_fn(params):
  batch_size = params['batch_size']

  images = tf.random.uniform(
      [total_examples, height, width, 3], minval=-1, maxval=1)

  dataset = tf.data.Dataset.from_tensor_slices(images)
  dataset = dataset.map(lambda images:{'image':images})

  dataset = dataset.batch(batch_size)
  return dataset

def model_fn(features, labels, params, mode):
   # Generate predictions, called 'output', from features['image']

  if mode == tf.estimator.ModeKeys.PREDICT:
    return tf.contrib.tpu.TPUEstimatorSpec(
        mode=mode,
        predictions={
            'predictions':output,
            'is_padding':features['is_padding']
        })

tpu_est = TPUEstimator(
    model_fn=model_fn,
    ...,
    predict_batch_size=16)

# Fully consume the generator so that TPUEstimator can shutdown the TPU
# system.
for item in tpu_est.predict(input_fn=input_fn):
  # Filter out item if the `is_padding` is 1.
  # Process the 'predictions'

导出

export_saved_model 导出 2 个元图，一个带有 saved_model.SERVING ，另一个带有 saved_model.SERVING 和 saved_model.TPU 标签。在服务时，这些标签用于选择要加载的适当元图。

在 TPU 上运行图之前，需要初始化 TPU 系统。如果使用 TensorFlow Serving model-server，则会自动完成。如果没有，请使用 session.run(tpu.initialize_system()) 。

API 有两个版本：1 或 2。

在 V1 中，导出的 CPU 图表原样为model_fn。导出的 TPU 图将 tpu.rewrite() 和 TPUPartitionedCallOp 包在 model_fn 周围，因此默认情况下 model_fn 在 TPU 上。要将操作放在 CPU 上，可以使用tpu.outside_compilation(host_call, logits)。

例子：

---

def model_fn(features, labels, mode, config, params):
  ...
  logits = ...
  export_outputs = {
    'logits':export_output_lib.PredictOutput(
      {'logits':logits})
  }

  def host_call(logits):
    class_ids = math_ops.argmax(logits)
    classes = string_ops.as_string(class_ids)
    export_outputs['classes'] =
      export_output_lib.ClassificationOutput(classes=classes)

  tpu.outside_compilation(host_call, logits)

  ...

在 V2 中，export_saved_model() 设置 params['use_tpu'] 标志以让用户知道代码是否正在导出到 TPU(或没有)。当 params['use_tpu'] 为 True 时，用户需要调用 tpu.rewrite() , TPUPartitionedCallOp 和/或 batch_function() 。

提示：推荐使用 V2，因为它更灵活(例如：批处理等)。