def testCopyToDeviceGpuWithMap(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") def generator(): for i in range(10): yield i, float(i), str(i) host_dataset = dataset_ops.Dataset.from_generator( generator, output_types=(dtypes.int32, dtypes.float32, dtypes.string)) device_dataset = host_dataset.apply( prefetching_ops.copy_to_device("/gpu:0")) def gpu_map_func(x, y, z): return math_ops.square(x), math_ops.square(y), z device_dataset = device_dataset.apply( prefetching_ops.map_on_gpu(gpu_map_func)) options = dataset_ops.Options() options.experimental_autotune = False device_dataset = device_dataset.with_options(options) with ops.device("/gpu:0"): iterator = device_dataset.make_initializable_iterator() next_element = iterator.get_next() with self.cached_session() as sess: sess.run(iterator.initializer) for i in range(10): x, y, z = sess.run(next_element) self.assertEqual(i**2, x) self.assertEqual(float(i**2), y) self.assertEqual(util_compat.as_bytes(str(i)), z) with self.assertRaises(errors.OutOfRangeError): sess.run(next_element)
def testInsideFunction(self): if test_util.is_gpu_available(): self.skipTest( "b/123899495: Colocation errors for critical sections in map on GPU" ) cs = critical_section_ops.CriticalSection() with ops.device( "/gpu:0" if test_util.is_gpu_available() else "/cpu:0"): v = resource_variable_ops.ResourceVariable(1) def fn(): return v.read_value() # map() creates a TensorFlow function. ds = dataset_ops.Dataset.range(1) if test_util.is_gpu_available(): ds = (ds.apply(prefetching_ops.copy_to_device("/gpu:0")).apply( prefetching_ops.map_on_gpu(lambda _: cs.execute(fn)))) else: ds = ds.map(lambda _: cs.execute(fn)) def get_first(): if context.executing_eagerly(): return self.evaluate( dataset_ops.make_one_shot_iterator(ds).get_next()) itr = dataset_ops.make_initializable_iterator(ds) self.evaluate([v.initializer, itr.initializer]) return self.evaluate(itr.get_next()) self.assertEqual(1, get_first())
def testInsideFunction(self): if test_util.is_gpu_available(): self.skipTest( "b/123899495: Colocation errors for critical sections in map on GPU") cs = critical_section_ops.CriticalSection() with ops.device("/gpu:0" if test_util.is_gpu_available() else "/cpu:0"): v = resource_variable_ops.ResourceVariable(1) def fn(): return v.read_value() # map() creates a TensorFlow function. ds = dataset_ops.Dataset.range(1) if test_util.is_gpu_available(): ds = (ds.apply(prefetching_ops.copy_to_device("/gpu:0")) .apply(prefetching_ops.map_on_gpu(lambda _: cs.execute(fn)))) else: ds = ds.map(lambda _: cs.execute(fn)) def get_first(): if context.executing_eagerly(): return self.evaluate(ds.make_one_shot_iterator().get_next()) itr = ds.make_initializable_iterator() self.evaluate([v.initializer, itr.initializer]) return self.evaluate(itr.get_next()) self.assertEqual(1, get_first())