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)

张量流中模型并行的实现

1 个答案: