parallel-processing - tensorflow 中模型并行性的实现

Question

我是tensorflow的初学者。我目前正在使用一个带有2个12GB GPU的系统。我想在两个GPU之间实现模型并行性，以训练大型模型。我一直在互联网，SO，tensorflow文档等中进行浏览，我能够找到模型并行性及其结果的解释，但在哪里找不到关于如何使用tensorflow实现的小型教程或小型代码片段。我的意思是我们必须在每一层交换交换激活信息，对不对？那么我们该怎么做呢？在 tensorflow 中是否有特定或较干净的方法来实现模型并行性？如果您可以向我建议一个可以学习实现它的地方，或者使用“MODEL PARALLELISM”在多个GPU上进行简单的mnist训练等简单代码，将对您有所帮助。

注意:我已经像CIFAR10中一样完成了数据并行性-multi gpu教程，但是我没有找到模型并行性的任何实现。

Answer 1

这是一个例子。该模型在GPU0上有一些部分，在GPU1上有一些部分，在CPU上有一些部分，因此这是三向模型并行性。

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)