def _create_ops(self, ds, ds_iterator, buffer_name, device0, device1):
ds_iterator_handle = ds_iterator.string_handle()
@function.defun(input_signature=[tensor_spec.TensorSpec([], dtypes.string)])
def _remote_fn(h):
remote_iterator = iterator_ops.Iterator.from_string_handle(
h, ds.output_types, ds.output_shapes)
return remote_iterator.get_next()
target = constant_op.constant(device0)
with ops.device(device1):
buffer_resource_handle = prefetching_ops.function_buffering_resource(
f=_remote_fn.get_concrete_function(),
output_types=[dtypes.float32],
target_device=target,
string_arg=ds_iterator_handle,
buffer_size=3,
shared_name=buffer_name)
with ops.device(device1):
prefetch_op = prefetching_ops.function_buffering_resource_get_next(
function_buffer_resource=buffer_resource_handle,
output_types=[dtypes.float32])
reset_op = prefetching_ops.function_buffering_resource_reset(
function_buffer_resource=buffer_resource_handle)
destroy_op = resource_variable_ops.destroy_resource_op(
buffer_resource_handle, ignore_lookup_error=True)
return (prefetch_op, reset_op, destroy_op)
def _create_ops(self, ds, ds_iterator, buffer_name, device0, device1):
ds_iterator_handle = ds_iterator.string_handle()
@function.defun(
input_signature=[tensor_spec.TensorSpec([], dtypes.string)])
def _remote_fn(h):
remote_iterator = iterator_ops.Iterator.from_string_handle(
h, ds.output_types, ds.output_shapes)
return remote_iterator.get_next()
target = constant_op.constant(device0)
with ops.device(device1):
buffer_resource_handle = prefetching_ops.function_buffering_resource(
f=_remote_fn.get_concrete_function(),
output_types=[dtypes.float32],
target_device=target,
string_arg=ds_iterator_handle,
buffer_size=3,
shared_name=buffer_name)
with ops.device(device1):
prefetch_op = prefetching_ops.function_buffering_resource_get_next(
function_buffer_resource=buffer_resource_handle,
output_types=[dtypes.float32])
reset_op = prefetching_ops.function_buffering_resource_reset(
function_buffer_resource=buffer_resource_handle)
destroy_op = resource_variable_ops.destroy_resource_op(
buffer_resource_handle, ignore_lookup_error=True)
return (prefetch_op, reset_op, destroy_op)
def testStringsGPU(self):
if not test_util.is_gpu_available():
self.skipTest("No GPU available")
device0 = "/job:localhost/replica:0/task:0/cpu:0"
device1 = "/job:localhost/replica:0/task:0/gpu:0"
ds = dataset_ops.Dataset.from_tensor_slices(["a", "b", "c"])
ds_iterator = ds.make_one_shot_iterator()
ds_iterator_handle = ds_iterator.string_handle()
@function.defun(
input_signature=[tensor_spec.TensorSpec([], dtypes.string)])
def _remote_fn(h):
remote_iterator = iterator_ops.Iterator.from_string_handle(
h, ds.output_types, ds.output_shapes)
return remote_iterator.get_next()
target = constant_op.constant(device0)
with ops.device(device1):
buffer_resource_handle = prefetching_ops.function_buffering_resource(
f=_remote_fn.get_concrete_function(),
output_types=[dtypes.string],
target_device=target,
string_arg=ds_iterator_handle,
buffer_size=3,
shared_name="strings")
with ops.device(device1):
prefetch_op = prefetching_ops.function_buffering_resource_get_next(
function_buffer_resource=buffer_resource_handle,
output_types=[dtypes.string])
destroy_op = resource_variable_ops.destroy_resource_op(
buffer_resource_handle, ignore_lookup_error=True)
with self.cached_session() as sess:
self.assertEqual([b"a"], self.evaluate(prefetch_op))
self.assertEqual([b"b"], self.evaluate(prefetch_op))
self.assertEqual([b"c"], self.evaluate(prefetch_op))
with self.assertRaises(errors.OutOfRangeError):
self.evaluate(prefetch_op)
self.evaluate(destroy_op)
def testStringsGPU(self):
if not test_util.is_gpu_available():
self.skipTest("No GPU available")
device0 = "/job:localhost/replica:0/task:0/cpu:0"
device1 = "/job:localhost/replica:0/task:0/gpu:0"
ds = dataset_ops.Dataset.from_tensor_slices(["a", "b", "c"])
ds_iterator = ds.make_one_shot_iterator()
ds_iterator_handle = ds_iterator.string_handle()
@function.defun(input_signature=[tensor_spec.TensorSpec([], dtypes.string)])
def _remote_fn(h):
remote_iterator = iterator_ops.Iterator.from_string_handle(
h, ds.output_types, ds.output_shapes)
return remote_iterator.get_next()
target = constant_op.constant(device0)
with ops.device(device1):
buffer_resource_handle = prefetching_ops.function_buffering_resource(
f=_remote_fn.get_concrete_function(),
output_types=[dtypes.string],
target_device=target,
string_arg=ds_iterator_handle,
buffer_size=3,
shared_name="strings")
with ops.device(device1):
prefetch_op = prefetching_ops.function_buffering_resource_get_next(
function_buffer_resource=buffer_resource_handle,
output_types=[dtypes.string])
destroy_op = resource_variable_ops.destroy_resource_op(
buffer_resource_handle, ignore_lookup_error=True)
with self.cached_session() as sess:
self.assertEqual([b"a"], self.evaluate(prefetch_op))
self.assertEqual([b"b"], self.evaluate(prefetch_op))
self.assertEqual([b"c"], self.evaluate(prefetch_op))
with self.assertRaises(errors.OutOfRangeError):
self.evaluate(prefetch_op)
self.evaluate(destroy_op)
def get_inputs_and_labels(function_buffering_resource, data_type):
"""Given a FunctionBufferingResource obtains images and labels from it."""
return prefetching_ops.function_buffering_resource_get_next(
function_buffer_resource=function_buffering_resource,
output_types=[data_type, tf.int32])