I try the new tensorflow function tf.contrib.data.prefetch to device.
My simple code example
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')
This gives the error:
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"]]
Is this new function not useable in combination with from_generator or is it something else?
