张量流 prefetch_to_device

Question

我尝试在设备上使用新的 tensorflow 函数 tf.contrib.data.prefetch。

我的简单代码示例

model = build_network()

N=1000

def gen():
    while True:
        batch = np.random.rand(N, 48, 48, 3)
        # Do some heavy calculation
        yield batch

dataset = tf.data.Dataset.from_generator(gen, tf.float32)
dataset = dataset.apply(tf.contrib.data.prefetch_to_device('/gpu:0'))

iterator = dataset.make_one_shot_iterator()
x = iterator.get_next()

output = model(x)

g = gen()
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for i in range(100):
        if i == 50:
            start = time.time()
        result = sess.run(output)
        #result = model.predict(next(g))
    end = time.time()
print('\nAverage time of one forward pass: {}\n'.format((end-start)/50))
print('Done')

这给出了错误：

InvalidArgumentError (see above for traceback): Cannot assign a device for operation 'IteratorGetDevice': Could not satisfy explicit device specification '/device:GPU:0' because no supported kernel for GPU devices is available. Colocation Debug Info: Colocation group had the following types and devices: IteratorToStringHandle: CPU IteratorGetDevice: CPU OneShotIterator: CPU

Colocation members and user-requested devices: OneShotIterator (OneShotIterator) IteratorGetDevice (IteratorGetDevice) /device:GPU:0 IteratorToStringHandle (IteratorToStringHandle)

Registered kernels: device='CPU'

[[Node: IteratorGetDevice = IteratorGetDevice_device="/device:GPU:0"]]

这个新功能不能与 from_generator 结合使用还是其他功能？

Answer 1

这是 TensorFlow 1.8rc0 候选版本中的错误。感谢您提请我们注意！

它现在已在 master branch and will be picked up in the next nightly build. I have also filed a cherry-pick to the 1.8 release branch 中修复，应该包含在 TensorFlow 1.8 的下一个候选版本（和最终版本）中。

张量流 prefetch_to_device

Tensorflow prefetch_to_device

python

tensorflow

tensorflow-datasets