Асинхронное вычисление в TensorFlow

Недавно я играл с TensorFlow, и я упомянул, что структура не может использовать все мои доступные вычислительные ресурсы. В разделе "Сверточные нейронные сети" упоминается, что

Наивное использование асинхронных обновлений параметров модели приводит к субоптимальным результатам обучения, поскольку индивидуальная модельная реплика может быть обучена устаревшей копии параметров модели. И наоборот, использование полностью синхронных обновлений будет таким же медленным, как и самая медленная реплика модели.

Хотя они упоминают это как в учебнике, так и в whitepaper, я действительно не нашел способ сделать асинхронную параллель вычисление на локальной машине. Возможно ли это? Или это часть распределенной выпущенной версии TensorFlow. Если это так, то как?

Ответ 1

Асинхронный градиентный спуск поддерживается в выпуске с открытым исходным кодом TensorFlow без изменения вашего графика. Самый простой способ сделать это - выполнить несколько параллельных шагов :

loss = ...

# Any of the optimizer classes can be used here.
train_op = tf.train.GradientDescentOptimizer(0.01).minimize(loss)

sess = tf.Session()
sess.run(tf.initialize_all_variables())

def train_function():
  # TODO: Better termination condition, e.g. using a `max_steps` counter.
  while True:
    sess.run(train_op)

# Create multiple threads to run `train_function()` in parallel
train_threads = []
for _ in range(NUM_CONCURRENT_STEPS):
  train_threads.append(threading.Thread(target=train_function))

# Start the threads, and block on their completion.
for t in train_threads:
  t.start()
for t in train_threads:
  t.join()

В этом примере настраиваются NUM_CONCURRENT_STEPS вызовы sess.run(train_op). Поскольку между этими потоками отсутствует координация, они выполняются асинхронно.

На самом деле более сложно добиться синхронного параллельного обучения (в настоящее время), поскольку для этого требуется дополнительная координация, чтобы гарантировать, что все реплики будут считывать одну и ту же версию параметров и что все их обновления станут видимыми одновременно. Пример multi-GPU для обучения CIFAR-10 выполняет синхронные обновления, делая несколько копий "башни" на графике обучения с общими параметрами и явно усредняя градиенты через башни перед применением обновления.


N.B. Код в этом ответе помещает все вычисления на одно и то же устройство, что не будет оптимальным, если на вашем компьютере будет несколько графических процессоров. Если вы хотите использовать все свои графические процессоры, следуйте примеру модели с несколькими GPU CIFAR-10 и создайте несколько "башен" с их операции, привязанные к каждому графическому процессору. Код выглядит примерно следующим образом:

train_ops = []

for i in range(NUM_GPUS):
  with tf.device("/gpu:%d" % i):
    # Define a tower on GPU `i`.
    loss = ...

    train_ops.append(tf.train.GradientDescentOptimizer(0.01).minimize(loss))

def train_function(train_op):
  # TODO: Better termination condition, e.g. using a `max_steps` counter.
  while True:
    sess.run(train_op)


# Create multiple threads to run `train_function()` in parallel
train_threads = []
for train_op in train_ops:
  train_threads.append(threading.Thread(target=train_function, args=(train_op,))


# Start the threads, and block on their completion.
for t in train_threads:
  t.start()
for t in train_threads:
  t.join()

Обратите внимание, что вам может быть удобно использовать "scope scope" для облегчения обмена информацией между башнями.