多GPU似乎无法在TensorFlow1.0上运行

时间:2017-03-09 17:12:38

标签: python machine-learning tensorflow

我正在使用TensorFlow 1.0,我开发了一个简单的程序来衡量性能。我有一个愚蠢的模型如下

def model(example_batch):
    h1 = tf.layers.dense(inputs=example_batch, units=64, activation=tf.nn.relu)
    h2 = tf.layers.dense(inputs=h1, units=2)
    return h2

和运行模拟的简单函数:

def testPerformanceFromMemory(model, iter=1000 num_cores=2):
  example_batch = tf.placeholder(np.float32, shape=(64, 128))
  for core in range(num_cores):
    with tf.device('/gpu:%d'%core):
      prediction = model(example_batch)
  init_op = tf.global_variables_initializer()
  sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
  sess.run(init_op)
  tf.train.start_queue_runners(sess=sess)
  input_array = np.random.random((64,128))
  for step in range(iter):
    myprediction = sess.run(prediction, feed_dict={example_batch:input_array})

如果我运行python脚本然后运行nvidia-smi命令我可以看到GPU0运行时使用率很高但GPU1的使用率为0%。

我读到了这个:https://www.tensorflow.org/tutorials/using_gpu而且:https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10/cifar10_multi_gpu_train.py但我不知道为什么我的例子不会在多个GPU中运行。

PS如果我从tensorflow存储库中执行ciphar 10示例,它将以multigpu模式运行。

编辑:正如mrry所说,我正在覆盖预测,所以我在这里发布正确的方法:

def testPerformanceFromMemory(model, iter=1000 num_cores=2):
  example_batch = tf.placeholder(np.float32, shape=(64, 128))
  prediction = []
  for core in range(num_cores):
    with tf.device('/gpu:%d'%core):
      prediction.append([model(example_batch)])
  init_op = tf.global_variables_initializer()
  sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
  sess.run(init_op)
  tf.train.start_queue_runners(sess=sess)
  input_array = np.random.random((64,128))
  for step in range(iter):
    myprediction = sess.run(prediction, feed_dict={example_batch:input_array})

1 个答案:

答案 0 :(得分:2)

查看您的程序,您将在不同的GPU设备上创建几个并行子图(通常称为“塔”),但在第一个prediction循环的每次迭代中覆盖for张量:

for core in range(num_cores):
  with tf.device('/gpu:%d'%core):
    prediction = model(example_batch)
# ...
for step in range(iter):
  myprediction = sess.run(prediction, feed_dict={example_batch:input_array})

因此,当您调用sess.run(prediction, ...)时,您将只运行在第一个for循环的最后一次迭代中创建的子图,该循环仅在一个GPU上运行。