tensorflow中模型并行的实现

Question

我目前正在开发一个具有 2 个 GPU，每个 12GB 的系统。我想在两个 GPU 上实现模型并行性来训练大型模型。我一直在浏览整个互联网、SO、tensorflow 文档等，我能够找到模型并行性及其结果的解释，但我在任何地方都找不到关于如何使用实现它的小教程或小代码 sn-ps张量流。我的意思是我们必须在每一层之后交换激活，对吗？那么我们该怎么做呢？是否有在 tensorflow 中实现模型并行性的特定或更简洁的方法？如果您能建议我一个可以学习实现它的地方或一个简单的代码（例如使用“模型并行”在多个 GPU 上进行 mnist 训练），那将非常有帮助。

注意：我已经完成了类似 CIFAR10 - 多 GPU 教程中的数据并行，但我还没有找到模型并行的任何实现。

Answer 1

这是一个例子。该模型在 GPU0 上有一些部分，在 GPU1 上有一些部分，在 CPU 上有一些部分，所以这是 3 路模型并行。

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)