def testIteratorOnDeviceGraphModeInitializableIterator(self):
if not test_util.is_gpu_available():
self.skipTest("No GPU available")
dataset = dataset_ops.Dataset.range(10)
dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0"))
iterator = dataset_ops.make_initializable_iterator(dataset)
data = iterator.get_next()
optional_data = iterator.get_next_as_optional()
with ops.colocate_with(dataset._variant_tensor):
dataset_device = test_ops.device_placement_op()
self.assertIn(b"GPU:0", self.evaluate(dataset_device))
with ops.colocate_with(iterator._iterator_resource):
iterator_device = test_ops.device_placement_op()
self.assertIn(b"GPU:0", self.evaluate(iterator_device))
with ops.colocate_with(data):
data_device = test_ops.device_placement_op()
self.assertIn(b"GPU:0", self.evaluate(data_device))
with ops.colocate_with(optional_data.get_value()):
get_value_device = test_ops.device_placement_op()
self.assertIn(b"GPU:0", self.evaluate(get_value_device))
with ops.colocate_with(optional_data.has_value()):
has_value_device = test_ops.device_placement_op()
self.assertIn(b"GPU:0", self.evaluate(has_value_device))
def testIteratorOnDeviceGraphModeOneShotIterator(self):
self.skipTest("TODO(b/169429285): tf.data.Dataset.make_one_shot_iterator "
"does not support GPU placement.")
dataset = dataset_ops.Dataset.range(10)
dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0"))
iterator = dataset_ops.make_one_shot_iterator(dataset)
data = iterator.get_next()
optional_data = iterator.get_next_as_optional()
with ops.colocate_with(dataset._variant_tensor):
dataset_device = test_ops.device_placement_op()
self.assertIn(b"GPU:0", self.evaluate(dataset_device))
with ops.colocate_with(iterator._iterator_resource):
iterator_device = test_ops.device_placement_op()
self.assertIn(b"GPU:0", self.evaluate(iterator_device))
with ops.colocate_with(data):
data_device = test_ops.device_placement_op()
self.assertIn(b"GPU:0", self.evaluate(data_device))
with ops.colocate_with(optional_data.get_value()):
get_value_device = test_ops.device_placement_op()
self.assertIn(b"GPU:0", self.evaluate(get_value_device))
with ops.colocate_with(optional_data.has_value()):
has_value_device = test_ops.device_placement_op()
self.assertIn(b"GPU:0", self.evaluate(has_value_device))
def func():
with ops.device('cpu:0'):
# Multiple CPU:0 devices match would be found, but the CPU:0 from the
# parent device scope should be picked.
x = test_ops.device_placement_op()
y = string_ops.string_upper(x)
packed_var_0 = array_ops.stack([x, y], 0)
return packed_var_0
def scan_fn(state, sample):
product = math_ops.matmul(sample, weights)
result.assign_add(product)
with ops.colocate_with(product):
device = test_ops.device_placement_op()
return state, device
def fun():
op = test_ops.device_placement_op()
return op
def fn2():
self.assertEqual("", constant_op.constant(3.0).op.device)
return test_ops.device_placement_op()
def fn2():
self.assertEqual("", constant_op.constant(3.0).op.device)
return test_ops.device_placement_op()
def fun(): op = test_ops.device_placement_op() return op