tf.estimator.tpu.TPUEstimator

Эстиматор с поддержкой TPU.

Наследуется от: Estimator

tf.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: ExportSavedModelApiVersion.V1 и V2.

В 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 для получения признаков и меток 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 с 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'}.

Параметр 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

Возвращает:

Вычисленные значения тензоров predictions.

train

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

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

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