TensorFlow:一个网络,两个GPU?

时间:2016-03-06 02:07:39

标签: python machine-learning neural-network tensorflow

我有一个卷积神经网络,有两个不同的输出流:

                         input
                           |
                         (...) <-- several convolutional layers
                           |
                       _________
    (several layers)   |       |    (several layers)
    fully-connected    |       |    fully-connected
    output stream 1 -> |       | <- output stream 2

我想计算/gpu:0上的流1和/gpu:1上的流2。不幸的是我无法正确设置它。

这次尝试:

...placeholders...
...conv layers...

with tf.device("/gpu:0"):
    ...stream 1 layers...
    nn_out_1 = tf.matmul(...)

with tf.device("/gpu:1"):
    ...stream 2 layers...
    nn_out_2 = tf.matmul(...)

运行速度慢(仅比1 GPU上的训练慢),有时会在输出中产生NaN值。我认为这可能是因为with语句可能无法正确同步。所以我添加了control_dependencies并将转化图层明确地放在/gpu:0上:

...placeholders...  # x -> input, y -> labels

with tf.device("/gpu:0"):
    with tf.control_dependencies([x, y]):
        ...conv layers...
        h_conv_flat = tf.reshape(h_conv_last, ...)

with tf.device("/gpu:0"):
    with tf.control_dependencies([h_conv_flat]):
        ...stream 1 layers...
        nn_out_1 = tf.matmul(...)

with tf.device("/gpu:1"):
    with tf.control_dependencies([h_conv_flat]):
        ...stream 2 layers...
        nn_out_2 = tf.matmul(...)

...但是通过这种方法,网络甚至都没有运行。无论我尝试过什么,它都抱怨输入没有被初始化:

tensorflow.python.framework.errors.InvalidArgumentError:
    You must feed a value for placeholder tensor 'x'
    with dtype float
    [[Node: x = Placeholder[dtype=DT_FLOAT, shape=[],
    _device="/job:localhost/replica:0/task:0/cpu:0"]()]]

如果没有with语句,网络只会在/gpu:0上进行培训并且运行良好 - 训练合理的东西,没有错误。

我做错了什么? TensorFlow无法将一个网络中不同的层流分割到不同的GPU吗?我总是必须拆分不同塔中的完整网络吗?

1 个答案:

答案 0 :(得分:2)

有一个如何在一个网络上使用多个gpus的示例 https://github.com/tensorflow/tensorflow/blob/master/tensorflow/models/image/cifar10/cifar10_multi_gpu_train.py 可能你可以复制代码。 也可以得到这样的东西

# Creates a graph.
c = []
for d in ['/gpu:2', '/gpu:3']:
with tf.device(d):
   a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3])
   b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2])
   c.append(tf.matmul(a, b))
with tf.device('/cpu:0'):
sum = tf.add_n(c)
# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# Runs the op.
print sess.run(sum)

查看:https://www.tensorflow.org/versions/r0.7/how_tos/using_gpu/index.html#using-multiple-gpus

最好的问候