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

Возвращаемое значение:

Оценённые значения тензоров predictions.

train

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

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

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