张量流中模型并行性的实现 [英] Implementation of model parallelism in tensorflow

问题描述

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

I am a beginner to tensorflow. I'm currently working on a system with 2 GPUs each of 12GB. I want to implement model parallelism across the two GPUs to train large models. I have been looking through all over the internet, SO, tensorflow documentation, etc, i was able to find the explanations of model parallelism and its results but nowhere did i find a small tutorial or small code snippets on how to implement it using tensorflow. I mean we have to exchange activations after every layer right? So how do we do that? Is there a specific or cleaner ways of implementing model parallelism in tensorflow? It would be very helpful if you could suggest me a place where i can learn to implement it or a simple code like mnist training on multiple GPU using 'MODEL PARALLELISM'.

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

Note: I have done data parallelism like in CIFAR10 - multi gpu tutorial but i haven't found any implementation of model parallelism.

推荐答案

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

Here's an example. The model has some parts on GPU0, some parts on GPU1 and some parts on CPU, so this is 3 way model parallelism.

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)

这篇关于张量流中模型并行性的实现的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！