tf.compat.v1.estimator.tpu.TPUEstimator

Эстиматор с поддержкой 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
)

TPU Estimator управляет собственной TensorFlow-графикой и сессией, поэтому он несовместим с поведением TF2. Мы рекомендуем перейти к новому tf.distribute.TPUStrategy. Подробнее см. в руководстве по TPU.

Описание

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

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

TPUEstimator преобразует глобальный размер пакета в params в размер пакета на фрагмент при вызове 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.

В версии 1 экспортированный CPU-график экспортируется без изменений. Экспортированный 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)

  ...

В версии 2 export_saved_model() устанавливает флаг params['use_tpu'], чтобы сообщить пользователю, экспортируется ли код на TPU (или нет). Когда params['use_tpu'] равно True, пользователям нужно вызвать tpu.rewrite(), TPUPartitionedCallOp и/или batch_function().

СОВЕТ: версия 2 рекомендуется, так как она более гибкая (например, пакетные операции и т. д.).

Методы

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 для получения признаков и меток Tensors. Далее, этот метод вызывает 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.ExportOutput, а входные данные всегда — приемники входных данных, предоставляемые 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 см. Использование формата SavedModel.

Этот метод строит новую графу, сначала вызывая serving_input_receiver_fn для получения признаков Tensor, а затем вызывая Estimator этого 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'}.

Параметр 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 этого 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 не работает. См.: issue/20506

train

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

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

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