Как запустить TensorFlow Estimator на нескольких графических процессорах с параллелизмом данных

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

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

У кого-нибудь есть хороший пример использования tf.учить.Оценщик? Или ссылку на учебник или что?

10
multi-gpu tensorflow tensorflow-gpu

3 ответов

Я думаю tf.ВНО.оценка.replicate_model_fn более чистое решение. Следующее из tf.ВНО.оценка.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(). Я следовал описанию и сделал модель squeezenet TPU на машине с 4 графическими процессорами. Мои модификации здесь.

5

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

один из способов запустить его data-parallel - это перебрать доступные устройства GPU и отправить куски вашего пакета в скопированные версии вашей модели (все сделано в вашем model_fn), а затем объединить результаты.

1

вы можете использовать область и устройство для этого:

 with tf.variable_scope(tf.get_variable_scope()):
  for i in xrange(FLAGS.num_gpus):
    with tf.device('/gpu:%d' % i):
      with tf.name_scope('%s_%d' % (cifar10.TOWER_NAME, i)) as scope:

полный пример есть: https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10/cifar10_multi_gpu_train.py

1