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