TensorFlow中模型并行化的实现

Question

我目前正在开发一个系统，该系统有 2 个 GPU，每个 12GB。我想在两个 GPU 之间实现模型并行性来训练大型模型。我一直在浏览整个互联网、SO、tensorflow 文档等，我能够找到模型并行性及其结果的解释，但我没有找到关于如何使用 tensorflow 实现它的小教程或小代码片段。我的意思是我们必须在每一层之后交换激活，对吗？那我们该怎么做呢？在 tensorflow 中是否有特定或更简洁的方法来实现模型并行性？如果你能建议我一个我可以学习实现它的地方或一个简单的代码，比如使用 'MODEL PARALLELISM'.

在多个 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)

TensorFlow中模型并行化的实现

Implementation of model parallelism in tensorflow

parallel-processing

distributed

tensorflow