Cálculo asíncrono en TensorFlow

Recientemente he estado jugando con TensorFlow y mencioné que el marco no puede usar todos los recursos computacionales disponibles. En el tutorial de Redes Neuronales Convolucionales mencionan que

El empleo ingenuo de actualizaciones asíncronas de los parámetros del modelo conduce a un rendimiento de entrenamiento por debajo del nivel óptimo, ya que una réplica individual del modelo podría recibir capacitación sobre una copia obsoleta de los parámetros del modelo. Por el contrario, emplear actualizaciones totalmente síncronas será tan lento como la réplica del modelo más lenta.

Aunque lo mencionaron tanto en el tutorial como en un documento técnico , realmente no encontré la forma de realizar el cálculo paralelo asíncrono en una máquina local. ¿Es incluso posible? O es parte de la versión distribuida para ser lanzada de TensorFlow. Si es así, entonces ¿cómo?

El descenso de gradiente asíncrono se admite en la versión de código abierto de TensorFlow, sin siquiera modificar su gráfico. La forma más sencilla de hacerlo es ejecutar varios pasos simultáneos en paralelo:

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, eg 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() 

Este ejemplo configura NUM_CONCURRENT_STEPS llamadas a sess.run(train_op) . Como no hay coordinación entre estos subprocesos, proceden de forma asíncrona.

En realidad, es más difícil lograr una capacitación paralela sincrónica (en la actualidad), ya que esto requiere una coordinación adicional para garantizar que todas las réplicas lean la misma versión de los parámetros y que todas sus actualizaciones se hagan visibles al mismo tiempo. El ejemplo multi-GPU para el entrenamiento CIFAR-10 realiza actualizaciones sincrónicas al hacer múltiples copias de la “torre” en el gráfico de entrenamiento con parámetros compartidos, y promediando explícitamente los gradientes en las torres antes de aplicar la actualización.


Nota: el código de esta respuesta coloca todos los cálculos en el mismo dispositivo, lo que no será óptimo si tiene varias GPU en su máquina. Si desea utilizar todas sus GPU, siga el ejemplo del modelo CIFAR-10 de múltiples GPU y cree múltiples “torres” con sus operaciones ancladas a cada GPU. El código se vería aproximadamente de la siguiente manera:

 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, eg 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() 

Tenga en cuenta que puede resultarle conveniente utilizar un “scope variable” para facilitar el intercambio de variables entre las torres.