Как запустить Оценщик Tensorflow на нескольких графических процессорах с данными parallelism

У меня есть стандартный оценочный калькулятор с некоторой моделью и хочу запускать его на нескольких графических процессорах вместо одного. Как это можно сделать с использованием данных parallelism?

Я искал Документы Tensorflow, но не нашел примера; только предложения, говорящие, что с Estimator будет легко.

Есть ли у кого-нибудь хороший пример с помощью tf.learn.Estimator? Или ссылку на учебник или так?

Ответ 1

Я думаю, что tf.contrib.estimator.replicate_model_fn является более чистым решением. Следующее из документации tf.contrib.estimator.replicate_model_fn,

...
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 'EstimatorSpec.train_op'.
    return EstimatorSpec(mode=mode, loss=loss,
                         train_op=optimizer.minimize(loss))

  #  No change for 'ModeKeys.EVAL' or 'ModeKeys.PREDICT'.
  return EstimatorSpec(...)
...
classifier = tf.estimator.Estimator(
  model_fn=tf.contrib.estimator.replicate_model_fn(model_fn))

Вам нужно обернуть оптимизатор с помощью tf.contrib.estimator.TowerOptimize и model_fn() с tf.contrib.estimator.replicate_model_fn(). Я последовал описанию и заставил модель TPU squeezenet работать на машине с 4 графическими процессорами. Мои модификации здесь.

Ответ 2

Стандартный пример: https://github.com/tensorflow/tensorflow/blob/r1.4/tensorflow/contrib/learn/python/learn/estimators/estimator.py

Одним из способов запуска параллельных данных было бы перебрать доступные графические устройства и отправить куски вашей партии в скопированные версии вашей модели (все сделано в вашей модели_fn), а затем объединить результаты.

Ответ 4

Я думаю, что это все, что вам нужно.

Ссылка: https://www.youtube.com/watch?v=bRMGoPqsn20

Подробнее: https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy

Объяснил: https://medium.com/tensorflow/multi-gpu-training-with-estimators-tf-keras-and-tf-data-ba584c3134db

NUM_GPUS = 8
dist_strategy = tf.contrib.distribute.MirroredStrategy(num_gpus=NUM_GPUS)
config = tf.estimator.RunConfig(train_distribute=dist_strategy)
estimator = tf.estimator.Estimator(model_fn,model_dir,config=config)

Ответ 5

Вы можете найти пример, используя tf.distribute.MirroredStrategy и tf.estimator.train_and_evaluate здесь.