Tensorflow-如何在GPU之间拆分批次以在经过训练的模型上进行预测?

时间:2019-03-04 16:03:45

标签: python tensorflow

我使用的不是创建但已修改的模型(来自此仓库https://github.com/GeorgeSeif/Semantic-Segmentation-Suite

我已经训练过模型,可以使用它们进行足够好的预测,但是我想遍历整个图像文件夹并在多个GPU之间分配工作。我不完全了解tf.device()的工作原理,而我尝试过的工作根本没有工作。

我认为我可以做这样的事情:

for i, d in enumerate(['\gpu:0', '\gpu:1']):
    with tf.device(d):
        output = sess.run(network, feed_dict={net_input: image_batch[i]})

但这实际上并没有将任务分配给不同的GPU,也没有引发错误。

我的问题是,是否可以在不显式修改网络代码预训练的情况下将不同的图像分配给单独GPU上的会话的不同实例。我想避免使用CUDA_VISIBLE_DEVICES = ...来运行两个不同的python脚本。

有一种简单的方法吗?

据我了解,操作的定义必须嵌套在“ with tf.device()”块中,但是,在推断操作时,只是模型和权重的加载,但是如果我将其放入“使用tf.device()“块时,我收到一条错误消息,指出该图形已经存在并且无法定义两次。

1 个答案:

答案 0 :(得分:0)

tf.device仅在构建图形而不执行图形时适用,因此将session.run调用包装在设备上下文中没有任何作用。

相反,我建议您使用tf replicator或tf分发策略(tf.distribute / tf.contrib.distribute取决于tf版本),尤其是MirroredStrategy。