Понимание распределения устройств, параллелизма (tf. while_loop) и функции tf.function в тензорном потоке

Question

Я пытаюсь понять параллелизм на GPU в тензорном потоке, так как мне нужно применить его к более уродливым графам.

import tensorflow as tf
from datetime import datetime

with tf.device('/device:GPU:0'):
    var = tf.Variable(tf.ones([100000], dtype=tf.dtypes.float32), dtype=tf.dtypes.float32)

@tf.function
def foo():
    return tf.while_loop(c, b, [i], parallel_iterations=1000)      #tweak

@tf.function
def b(i):
    var.assign(tf.tensor_scatter_nd_update(var, tf.reshape(i, [-1,1]), tf.constant([0], dtype=tf.dtypes.float32)))
    return tf.add(i,1)

with tf.device('/device:GPU:0'):
    i = tf.constant(0)
    c = lambda i: tf.less(i,100000)

start = datetime.today()
with tf.device('/device:GPU:0'):
    foo()
print(datetime.today()-start)

В приведенном выше коде var - это тензор с длиной 100000, элементы которого обновляютсякак показано выше.Когда я изменяю значения parallel_iterations с 10, 100, 1000, 10000. Вряд ли есть какая-либо разница во времени (все на 9,8 с), даже если явно упоминается переменная parallel_iterations.

Я хочу, чтобы это происходило параллельно на GPU.Как я могу это реализовать?

Answer 1

Одним из методов является использование стратегии и области распространения:

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
  inputs = tf.keras.layers.Input(shape=(1,))
  predictions = tf.keras.layers.Dense(1)(inputs)
  model = tf.keras.models.Model(inputs=inputs, outputs=predictions)
  model.compile(loss='mse',
                optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.2))

Другой вариант - дублирование операций на каждом устройстве:

# Replicate your computation on multiple GPUs
c = []
for d in ['/device:GPU:2', '/device:GPU:3']:
  with tf.device(d):
    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3])
    b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2])
    c.append(tf.matmul(a, b))
with tf.device('/cpu:0'):
  sum = tf.add_n(c)

См. это руководство для более подробной информации