tf.compat.v1.estimator.tpu.TPUEstimator

Estimator с поддержкой TPU.

Наследуется от: 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
)

TPUEstimator также поддерживает обучение на CPU и GPU. Вам не нужно определять отдельный tf.estimator.Estimator.

TPUEstimator обрабатывает многие детали выполнения на устройствах TPU, такие как дублирование входных данных и моделей для каждого ядра и периодическое возвращение на хост для выполнения хуков.

TPUEstimator преобразует глобальный размер пакета в параметрах в размер пакета на фрагмент при вызове input_fn и model_fn. Пользователи должны указать глобальный размер пакета в конструкторе, а затем получить размер пакета для каждого фрагмента в input_fn и model_fn с помощью params['batch_size'].

• Для обучения model_fn получает размер пакета на ядро; input_fn может получить размер пакета на ядро или на хост в зависимости от per_host_input_for_training в TPUConfig (подробности см. в документации к TPUConfig).

• Для оценки и предсказания model_fn получает размер пакета на ядро, а input_fn получает размер пакета на хост.

Оценка

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 и возвращает словарь, где имя метрики соответствует результату вызова функции метрики, а именно кортежу (metric_tensor, update_op).

Можно установить use_tpu в False для тестирования. Все обучение, оценка и предсказание будут выполняться на CPU. input_fn и model_fn получат train_batch_size или eval_batch_size без изменений как params['batch_size'].

Текущие ограничения:

1. Оценка на TPU работает только на одном хосте (один рабочий TPU), за исключением режима BROADCAST.

2. input_fn для оценки не должно генерировать исключение конца входных данных (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 — это экспериментальная функция для поддержки предсказания с большими пакетами. Она не предназначена для систем, критичных к задержке. Кроме того, из-за некоторых проблем с удобством использования для предсказания с небольшим набором данных, использование CPU .predict, т. е. создание нового экземпляра TPUEstimator с use_tpu=False, может быть удобнее.

Примечание: В отличие от обучения/оценки на TPU, input_fn для предсказания должно генерировать исключение конца входных данных (OutOfRangeError или StopIteration). Это служит сигналом остановки для TPUEstimator. Точнее, операции, созданные input_fn, генерируют один пакет данных. API predict() обрабатывает один пакет за раз. При достижении конца источника данных должно быть сгенерировано исключение конца входных данных одной из этих операций. Обычно пользователь не должен делать это вручную. Пока набор данных не повторяется бесконечно, API tf.data автоматически сгенерирует исключение конца входных данных после того, как будет сгенерирован последний пакет.

Примечание: 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(). В качестве альтернативы можно использовать inference_on_tpu(), который представляет собой удобную оболочку для этих трёх функций.

def model_fn(features, labels, mode, config, params):
  ...
  # This could be some pre-processing on CPU like calls to input layer with
  # embedding columns.
  x2 = features['x'] * 2

  def computation(input_tensor):
    return layers.dense(
        input_tensor, 1, kernel_initializer=init_ops.zeros_initializer())

  inputs = [x2]
  if params['use_tpu']:
    predictions = array_ops.identity(
        tpu_estimator.inference_on_tpu(computation, inputs,
        num_batch_threads=1, max_batch_size=2, batch_timeout_micros=100),
        name='predictions')
  else:
    predictions = array_ops.identity(
        computation(*inputs), name='predictions')
  key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
  export_outputs = {
      key: export_lib.PredictOutput({'prediction': predictions})
  }
  ...

СОВЕТ: Рекомендуется использовать V2, так как он более гибкий (например, пакетная обработка и т. д.).

Методы

eval_dir

Просмотреть исходный код

eval_dir(
    name=None
)

Отображает имя каталога, куда сохраняются метрики оценки.

evaluate

Просмотреть исходный код

evaluate(
    input_fn, steps=None, hooks=None, checkpoint_path=None, name=None
)

Оценивает модель на основе данных оценки input_fn.

Для каждого шага вызывается input_fn, возвращающая одну партию данных. Оценка выполняется до:

• steps партий будут обработаны, или
• input_fn возникает исключение конца ввода (tf.errors.OutOfRangeError или StopIteration).

experimental_export_all_saved_models

Просмотреть исходный код

experimental_export_all_saved_models(
    export_dir_base, input_receiver_fn_map, assets_extra=None, as_text=False,
    checkpoint_path=None
)

Экспортирует SavedModel с tf.MetaGraphDefs для каждого запрошенного режима.

Для каждого режима, переданного через input_receiver_fn_map, этот метод создаёт новую графу, вызывая input_receiver_fn для получения признаков и меток Tensor. Затем этот метод вызывает Estimator's model_fn в переданном режиме для генерации графа модели на основе этих признаков и меток, и восстанавливает заданный контрольный пункт (или, в случае его отсутствия, последний контрольный пункт) в графе. Только один из режимов используется для сохранения переменных в SavedModel (порядок предпочтения: tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL, затем tf.estimator.ModeKeys.PREDICT), таким образом, до трёх tf.MetaGraphDefs сохраняются с одним набором переменных в одном каталоге SavedModel.

Для переменных и tf.MetaGraphDefs, создаётся каталог экспорта с отметкой времени ниже export_dir_base, и в нём записывается SavedModel с содержанием tf.MetaGraphDef для данного режима и его связанных сигнатур.

Для предсказания, экспортированный MetaGraphDef предоставит по одному SignatureDef для каждого элемента словаря export_outputs, возвращённого из model_fn, с именами, использующими те же ключи. Один из этих ключей всегда tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, указывающий на то, какая подпись будет использоваться, когда запрос на обслуживание не указывает её. Для каждой подписи выходы предоставляются соответствующими tf.estimator.export.ExportOutputs, а входные данные всегда — приёмники входных данных, предоставляемые serving_input_receiver_fn.

Для обучения и оценки, train_op хранится в дополнительном наборе, а потеря, метрики и предсказания включаются в SignatureDef для рассматриваемого режима.

Дополнительные ресурсы могут быть записаны в SavedModel через аргумент assets_extra. Это должен быть словарь, где каждый ключ указывает путь назначения (включая имя файла) относительно каталога assets.extra. Соответствующее значение даёт полный путь исходного файла, который нужно скопировать. Например, простой случай копирования одного файла без переименования задаётся как {'my_asset_file.txt': '/path/to/my_asset_file.txt'}.

export_saved_model

Просмотреть исходный код

export_saved_model(
    export_dir_base, serving_input_receiver_fn, assets_extra=None, as_text=False,
    checkpoint_path=None, experimental_mode=ModeKeys.PREDICT
)

Экспортирует граф инференции в виде SavedModel в заданный каталог.

Для подробного руководства см. SavedModel от Estimators.

Этот метод строит новую графу, сначала вызвав serving_input_receiver_fn для получения признаков Tensor, а затем вызвав Estimator's model_fn для генерации графа модели на основе этих признаков. Он восстанавливает заданный контрольный пункт (или, в случае его отсутствия, последний контрольный пункт) в этой графе в новой сессии. Наконец, он создаёт каталог экспорта с отметкой времени ниже заданного export_dir_base, и записывает в него SavedModel с одной tf.MetaGraphDef сохранённой из этой сессии.

Экспортированный MetaGraphDef предоставит по одному SignatureDef для каждого элемента словаря export_outputs, возвращённого из model_fn, используя те же ключи. Один из этих ключей всегда tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, указывающий на то, какая подпись будет использоваться, когда запрос на обслуживание не указывает её. Для каждой подписи выходы предоставляются соответствующими tf.estimator.export.ExportOutputs, а входные данные всегда — приёмники входных данных, предоставляемые serving_input_receiver_fn.

Дополнительные ресурсы могут быть записаны в SavedModel через аргумент assets_extra. Это должен быть словарь, где каждый ключ указывает путь назначения (включая имя файла) относительно каталога assets.extra. Соответствующее значение указывает полный путь исходного файла, который нужно скопировать. Например, простой случай копирования одного файла без переименования задаётся как {'my_asset_file.txt': '/path/to/my_asset_file.txt'}.

Параметр experimental_mode может использоваться для экспорта единственной графы обучения/валидации/предсказания в формате SavedModel. Полные документы см. в experimental_export_all_saved_models.

export_savedmodel

Просмотреть исходный код

export_savedmodel(
    export_dir_base, serving_input_receiver_fn, assets_extra=None, as_text=False,
    checkpoint_path=None, strip_default_attrs=False
)

Экспортирует графу вывода в формате SavedModel в указанный каталог. (устаревшая функция)

Подробное руководство см. в SavedModel от Estimators.

Этот метод строит новую графу, сначала вызывая serving_input_receiver_fn для получения функций Tensor, а затем вызывая Estimator's model_fn для создания графы модели на основе этих функций. Он восстанавливает заданную контрольную точку (или, если её нет, последнюю контрольную точку) в эту графу в новой сессии. Наконец, он создаёт каталог экспорта с отметкой времени ниже заданного export_dir_base, и записывает в него SavedModel содержащую единственный tf.MetaGraphDef сохранённый из этой сессии.

Экспортированный MetaGraphDef предоставит по одному SignatureDef для каждого элемента словаря export_outputs, возвращённого model_fn, используя те же ключи. Один из этих ключей всегда tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, указывая, какая сигнатура будет обработана, если в запросе сервера не указана. Для каждой сигнатуры выходные данные предоставляются соответствующими tf.estimator.export.ExportOutput, а входные данные всегда являются приемниками ввода, предоставленными serving_input_receiver_fn.

Дополнительные ресурсы могут быть записаны в SavedModel через аргумент assets_extra. Это должен быть словарь, где каждый ключ указывает путь назначения (включая имя файла) относительно каталога assets.extra. Соответствующее значение указывает полный путь исходного файла, который нужно скопировать. Например, простой случай копирования одного файла без переименования задаётся как {'my_asset_file.txt': '/path/to/my_asset_file.txt'}.

get_variable_names

Просмотреть исходный код

get_variable_names()

Возвращает список всех имён переменных в этой модели.

get_variable_value

Просмотреть исходный код

get_variable_value(
    name
)

Возвращает значение переменной по заданному имени.

latest_checkpoint

Просмотреть исходный код

latest_checkpoint()

Находит имя файла последней сохранённой контрольной точки в model_dir.

predict

Просмотреть исходный код

predict(
    input_fn, predict_keys=None, hooks=None, checkpoint_path=None,
    yield_single_examples=True
)

Возвращает предсказания для заданных признаков.

Обратите внимание, что интерлейсинг двух выходов predict не работает. См.: вопрос/20506

train

Просмотреть исходный код

train(
    input_fn, hooks=None, steps=None, max_steps=None, saving_listeners=None
)

Обучает модель с использованием обучающих данных input_fn.