def testTensorsExplicitPrefetchToDevice(self): ds = Dataset.from_tensor_slices([0., 1.]) ds = ds.apply(prefetching_ops.prefetch_to_device(test.gpu_device_name())) with self.assertRaisesRegexp(TypeError, 'prefetch_to_device'): datasets.Iterator(ds) for i, x in enumerate(ds): with ops.device(test.gpu_device_name()): x = math_ops.add(x, x) self.assertEqual(float(i) + float(i), x.numpy())
def testTensorsPlacedOnDevice(self): ds = Dataset.from_tensors([0., 1.]) with ops.device(test.gpu_device_name()): x = datasets.Iterator(ds).next() x = math_ops.add(x, x) self.assertAllEqual([0., 2.], x.numpy())
def testGpuDefinedDataset(self): with ops.device(test.gpu_device_name()): ds = Dataset.from_tensors([0., 1.]) for x in ds: y = math_ops.add(x, x) self.assertAllEqual([0., 2.], y.numpy())
def testGpuTensor(self): ds = dataset_ops.Dataset.from_tensors([0., 1.]) with ops.device(test.gpu_device_name()): for x in ds: y = math_ops.add(x, x) self.assertAllEqual([0., 2.], y.numpy())