支持 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，因為它更靈活(例如：批處理等)。