tf.contrib.estimator.replicate_model_fn

Репликация Estimator.model_fn по GPU. (устаревшая)

tf.contrib.estimator.replicate_model_fn(
    model_fn, loss_reduction=losses.Reduction.SUM_BY_NONZERO_WEIGHTS, devices=None
)

Указанная model_fn определяет одиночный прямой проход модели. Для репликации такой модели по GPU каждый GPU получает свою копию прямого прохода (также известный как башня). Входные признаки и метки разбиваются на части, соответствующие количеству GPU. Каждая башня вычисляет потерю на основе своего входного значения. Для каждой такой потери вычисляются градиенты. После этого доступные потери агрегируются, образуя агрегированную потерю. Доступные градиенты суммируются. Затем они обновляют веса, используя указанный оптимизатор.

Если devices являются None, тогда все доступные GPU будут использоваться для репликации. Если GPU недоступны, тогда модель будет размещена на CPU.

Поддерживаются два режима локальной репликации по доступным GPU:

1) Если обнаружен ровно 1 GPU, то переменные и операции размещаются на GPU. 2) Если обнаружено более 1 GPU, то переменные будут размещены на CPU. Копии операций размещаются на каждом отдельном GPU.

Вот пример того, как можно использовать свою model_fn для работы по GPU:

...
def model_fn(...):  # See `model_fn` in `Estimator`.
  loss = ...
  optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
  optimizer = tf.contrib.estimator.TowerOptimizer(optimizer)
  if mode == tf.estimator.ModeKeys.TRAIN:
    #  See the section below on <a href="/api_docs/python/tf/estimator/EstimatorSpec.md#train_op"><code>EstimatorSpec.train_op</code></a>.
    return EstimatorSpec(mode=mode, loss=loss,
                         train_op=optimizer.minimize(loss))

  #  No change for <a href="/api_docs/python/tf/estimator/ModeKeys.md#EVAL"><code>ModeKeys.EVAL</code></a> or <a href="/api_docs/python/tf/estimator/ModeKeys.md#PREDICT"><code>ModeKeys.PREDICT</code></a>.
  return EstimatorSpec(...)
...
classifier = tf.estimator.Estimator(
  model_fn=tf.contrib.estimator.replicate_model_fn(model_fn))

См. DNNClassifierIntegrationTest для примера с готовым оценщиком.

На EstimatorSpec.train_op: model_fn возвращает EstimatorSpec.train_op для tf.estimator.GraphKeys.TRAIN. Обычно он выводится с помощью оптимизатора. Ожидается, что башни будут заполнять его аналогичным образом. Градиенты со всех башен сокращаются и применяются в последней башне. Для достижения этого в случае нескольких башен необходимо использовать TowerOptimizer. См. TowerOptimizer.

При фрагментации входных признаков и меток: Входные признаки и метки разбиваются для потребления каждой башней. Они разбиваются по размеру 0. Признаки и метки должны быть в формате пакет-главный.

По алгоритмам сокращения: Были выбраны определенные алгоритмы для агрегирования результатов вычислений на нескольких башнях:

• Потери со всех башен сокращаются в соответствии с loss_reduction.
• Градиенты со всех башен сокращаются в соответствии с loss_reduction для каждой обучаемой переменной.
• eval_metrics_ops сокращаются по метрике с помощью reduce_mean.
• EstimatorSpec.predictions и EstimatorSpec.export_outputs сокращаются с помощью конкатенации.
• Для всех остальных полей EstimatorSpec берутся значения первой башни.

По распределению переменных: Переменные не дублируются между башнями. Вместо этого они размещаются на единственном устройстве, как определено выше, и совместно используются во всех башнях.

Накладные расходы:

Если указано только одно устройство, то агрегация потерь и градиентов не происходит. Репликация заключается в размещении model_fn на указанном устройстве.

О текущих ограничениях:

• predictions не поддерживаются для ModeKeys.EVAL. Они необходимы для tf.contrib.estimator.add_metrics.