python之Tensorflow prefetch_to_device

我尝试使用新的 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 结合使用还是其他？

请您参考如下方法:

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

它现在固定在 master branch 中并将在下一个夜间构建中被拾取。我还提交了 cherry-pick to the 1.8 release branch ，它应该包含在 TensorFlow 1.8 的下一个候选版本(和最终版本)中。