Реализация параллелизма модели в тензорном потоке

Question

Я новичок в тензорном потоке. В настоящее время я работаю над системой с 2 ​​графическими процессорами, каждый из которых составляет 12 ГБ. Я хочу реализовать параллелизм модели на двух графических процессорах для подготовки больших моделей. Я просматривал всю информацию в Интернете, SO, документацию о тензорном потоке и т. Д., Я смог найти объяснения параллелизма модели и ее результатов, но нигде я не нашел небольшой учебник или небольшие фрагменты кода о том, как его реализовать, используя tenorflow. Я имею в виду, что мы должны обмениваться активациями сразу после каждого слоя? Итак, как мы это делаем? Существуют ли конкретные или более чистые способы реализации параллелизма модели в тензорном потоке? Было бы очень полезно, если бы вы предложили мне место, где я могу научиться его реализовывать, или простой код, например, обучение mnist на нескольких графических процессорах, используя «MODEL PARALLELISM».

Примечание: Я выполнил параллелизм данных, как в CIFAR10 - многопроцессорный учебник gpu, но я не нашел реализации параллелизма модели.

Answer 1

Вот пример. Модель имеет некоторые части на GPU0, некоторые части на GPU1 и некоторые части на процессоре, поэтому это трехходовой параллелизм модели.

with tf.device("/gpu:0"): 
    a = tf.Variable(tf.ones(())) 
    a = tf.square(a) 
with tf.device("/gpu:1"): 
    b = tf.Variable(tf.ones(())) 
    b = tf.square(b) 
with tf.device("/cpu:0"): 
    loss = a+b 
opt = tf.train.GradientDescentOptimizer(learning_rate=0.1) 
train_op = opt.minimize(loss) 

sess = tf.Session() 
sess.run(tf.global_variables_initializer()) 
for i in range(10): 
    loss0, _ = sess.run([loss, train_op]) 
    print("loss", loss0)