Tensorflow 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 (回溯参见上文)：无法为操作“IteratorGetDevice”分配设备:无法满足显式设备规范/设备: GPU :0，因为GPU设备没有支持的内核。协同定位调试信息:协同定位组有以下类型和设备: IteratorToStringHandle: CPU IteratorGetDevice: CPU OneShotIterator: CPU 定位成员和用户请求的设备: OneShotIterator (OneShotIterator) IteratorGetDevice (IteratorGetDevice) /device:GPU:0 IteratorToStringHandle (IteratorToStringHandle) 注册内核:设备=‘CPU’ [节点: IteratorGetDevice = IteratorGetDevice_device="/device:GPU:0"]

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

Answer 1

这是TensorFlow 1.8rc0版本候选版本中的一个bug。谢谢你把它带给我们的注意！

它现在在主支中被修复，并将在下一个晚上的构建中得到。我还提交了一个樱桃-选择1.8发布分支，它应该包含在TensorFlow 1.8的下一个发布候选版本(和最终版本)中。