Question

当我使用此代码时，tensorflow不会显示我的tf.add运算符的设备放置位置。如何知道此运算符是在GPU还是CPU上完成的？

import tensorflow as tf

c_0 = tf.constant(0, name="c")

with tf.device("/device:CPU:0"):
    c_1 = tf.constant(51, name="c")

with tf.device("/device:GPU:0"):
    c_2 = tf.constant(56, name="c")


def add_2(val):
    return tf.add(val, 2)


dataset = tf.data.Dataset.from_tensor_slices([c_0, c_1, c_2])
dataset = dataset.map(add_2)
iterator = dataset.make_one_shot_iterator()
get_next = iterator.get_next()

with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
    print(sess.run(get_next))
    print(sess.run(get_next))
    print(sess.run(get_next))

输出为：

tensors/component_0: (Pack): /job:localhost/replica:0/task:0/device:GPU:0
IteratorToStringHandle: (IteratorToStringHandle): /job:localhost/replica:0/task:0/device:CPU:0
IteratorGetNext: (IteratorGetNext): /job:localhost/replica:0/task:0/device:CPU:0
c: (Const): /job:localhost/replica:0/task:0/device:GPU:0
c_1: (Const): /job:localhost/replica:0/task:0/device:CPU:0
c_2: (Const): /job:localhost/replica:0/task:0/device:GPU:0
OneShotIterator: (OneShotIterator): /job:localhost/replica:0/task:0/device:CPU:0
2
53
58

Answer 1

默认情况下，Dataset API在CPU上运行。它不会记录数据集对象内的操作。但是，您可以将迭代器显式映射到GPU。

dataset = tf.data.Dataset.from_tensor_slices([c_0, c_1, c_2])
dataset = dataset.map(add_2)
dataset = dataset.apply(tf.data.experimental.copy_to_device("/gpu:0")).prefetch(1)

这仅适用于make_initializable_iterator() 看看这个链接

https://github.com/tensorflow/tensorflow/issues/19244

https://github.com/tensorflow/tensorflow/issues/13610

如何获取张量流日志设备在数据集上的位置？

1 个答案: