def test_generated_rpc_ops(self): @eager_def_function.function(input_signature=[ tensor_spec.TensorSpec([], dtypes.int32), tensor_spec.TensorSpec([], dtypes.int32) ]) def remote_fn(a, b): return math_ops.multiply(a, b) concrete_remote_fn = remote_fn.get_concrete_function() a = variables.Variable(2, dtype=dtypes.int32) b = variables.Variable(3, dtype=dtypes.int32) port = portpicker.pick_unused_port() address = "localhost:{}".format(port) server_resource = rpc_ops.gen_rpc_ops.rpc_server( server_address=address) rpc_ops.gen_rpc_ops.rpc_server_register( server_resource, f=concrete_remote_fn, captured_inputs=concrete_remote_fn.captured_inputs, output_specs=rpc_ops.get_output_specs_from_function( concrete_remote_fn), method_name="multiply") rpc_ops.gen_rpc_ops.rpc_server_start(server_resource) client_handle, _ = rpc_ops.gen_rpc_ops.rpc_client( server_address=address, timeout_in_ms=5000) future_resource, deleter = rpc_ops.gen_rpc_ops.rpc_call( client_handle, args=[a, b], method_name="multiply", timeout_in_ms=0) error_code, _ = rpc_ops.gen_rpc_ops.rpc_check_status(future_resource) self.assertAllEqual(error_code, 0) self.assertAllEqual( rpc_ops.gen_rpc_ops.rpc_get_value(future_resource, Tout=[dtypes.int32]), [6]) resource_variable_ops.EagerResourceDeleter( handle=server_resource, handle_device=server_resource.device) resource_variable_ops.EagerResourceDeleter( handle=client_handle, handle_device=client_handle.device) rpc_ops.gen_rpc_ops.delete_rpc_future_resource(future_resource, deleter)
def __init__(self, num_threads, display_name=None, max_intra_op_parallelism=1): """Creates a `PrivateThreadPool` with the given number of threads.""" if context.executing_eagerly(): shared_name = _generate_shared_name("privatethreadpool") if compat.forward_compatible(2019, 8, 3): self._resource = ged_ops.thread_pool_handle( num_threads=num_threads, max_intra_op_parallelism=max_intra_op_parallelism, display_name=display_name, shared_name=shared_name) else: self._resource = ged_ops.experimental_thread_pool_handle( num_threads=num_threads, max_intra_op_parallelism=max_intra_op_parallelism, display_name=display_name, shared_name=shared_name) self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device=context.context().device_name) else: if compat.forward_compatible(2019, 8, 3): self._resource = ged_ops.thread_pool_handle( num_threads=num_threads, max_intra_op_parallelism=max_intra_op_parallelism, display_name=display_name) else: self._resource = ged_ops.experimental_thread_pool_handle( num_threads=num_threads, max_intra_op_parallelism=max_intra_op_parallelism, display_name=display_name)
def __init__(self, address: str, name: str = "", list_registered_methods=False, timeout_in_ms=0): self._client_handle, methods = gen_rpc_ops.rpc_client( shared_name=name, server_address=address, list_registered_methods=list_registered_methods, timeout_in_ms=timeout_in_ms) if context.executing_eagerly(): self._handle_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._client_handle, handle_device=self._client_handle.device) else: raise NotImplementedError( "Client creation is supported only in eager mode.") self._server_address = address decoder = nested_structure_coder.StructureCoder() self._method_registry = {} for method in methods.numpy(): m = rpc_pb2.RegisteredMethod() m.ParseFromString(method) output_specs = decoder.decode_proto(m.output_specs) input_specs = decoder.decode_proto(m.input_specs) self._method_registry[m.method] = output_specs # TODO(ishark): Perhaps doc string can also be taken as input during # function registration. doc_string = "RPC Call for " + m.method + " method to server " + address self._add_method(m.method, output_specs, input_specs, self._client_handle, doc_string)
def __init__(self, dataset): """Creates a new iterator over the given dataset. For example: ```python dataset = tf.data.Dataset.range(4) for x in Iterator(dataset): print(x) ``` Tensors produced will be placed on the device on which this iterator object was created. Args: dataset: A `tf.data.Dataset` object. Raises: TypeError: If `dataset` is an unsupported type. RuntimeError: When invoked without eager execution enabled. """ if isinstance(dataset, prefetching_ops._PrefetchToDeviceDataset): # pylint: disable=protected-access raise TypeError( "`tf.contrib.data.prefetch_to_device()` is not compatible with " "`tf.contrib.eager.Iterator`. Use `for ... in dataset:` to iterate " "over the dataset instead.") super(Iterator, self).__init__(dataset) if not context.context().device_spec.device_type: is_remote_device = False else: is_remote_device = context.context( ).device_spec.device_type != "CPU" self._buffer_resource_handle = None if is_remote_device: with ops.device("/device:CPU:0"): iter_string_handle = gen_dataset_ops.iterator_to_string_handle( self._resource) @function.Defun(dtypes.string) def remote_fn(h): remote_iterator = iterator_ops.Iterator.from_string_handle( h, self.output_types, self.output_shapes, self.output_classes) return remote_iterator.get_next() remote_fn.add_to_graph(None) target = constant_op.constant("/device:CPU:0") with ops.device(self._device): self._buffer_resource_handle = prefetching_ops.function_buffering_resource( # pylint: disable=line-too-long string_arg=iter_string_handle, output_types=self._flat_output_types, f=remote_fn, target_device=target, buffer_size=10, container="", shared_name=_generate_shared_name( "contrib_eager_iterator_function_buffer_resource")) self._buffer_resource_deleter = resource_variable_ops.EagerResourceDeleter( # pylint: disable=line-too-long handle=self._buffer_resource_handle, handle_device=self._device)
def __init__(self, address: str): self._server_handle = gen_rpc_ops.rpc_server(address) if context.executing_eagerly(): self._handle_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._server_handle, handle_device=self._server_handle.device) else: raise NotImplementedError("Please create the server outside tf.function.")
def __init__(self, resource, init_op_fn): self._resource = resource # TODO(nickfelt): cache constructed ops in graph mode self._init_op_fn = init_op_fn if context.executing_eagerly() and self._resource is not None: self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device="cpu:0")
def __init__(self, resource, init_op_fn): self._resource = resource self._init_op_fn = init_op_fn init_op = self.init() if context.executing_eagerly(): self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device="cpu:0") else: ops.add_to_collection(_SUMMARY_WRITER_INIT_COLLECTION_NAME, init_op)
def __init__(self, shared_name, init_op_fn, name=None): self._resource = gen_summary_ops.summary_writer( shared_name=shared_name, name=name) self._init_op_fn = init_op_fn self._init_op = init_op_fn(self._resource) self._closed = False if context.executing_eagerly(): self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device="cpu:0") else: ops.add_to_collection(_SUMMARY_WRITER_INIT_COLLECTION_NAME, self._init_op)
def __init__(self, num_threads, display_name=None): """Creates a `PrivateThreadPool` with the given number of threads.""" if context.executing_eagerly(): shared_name = _generate_shared_name("privatethreadpool") self._resource = gen_dataset_ops.thread_pool_handle( num_threads=num_threads, display_name=display_name, shared_name=shared_name) self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device=context.context().device_name) else: self._resource = gen_dataset_ops.thread_pool_handle( num_threads=num_threads, display_name=display_name)
def __init__(self, shared_name, init_op_fn, name=None, v2=False): self._resource = gen_summary_ops.summary_writer( shared_name=shared_name, name=name) # TODO(nickfelt): cache other constructed ops in graph mode self._init_op_fn = init_op_fn self._init_op = init_op_fn(self._resource) self._v2 = v2 self._closed = False if context.executing_eagerly(): self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device="cpu:0") else: global _SUMMARY_WRITER_INIT_OP key = ops.get_default_graph()._graph_key # pylint: disable=protected-access _SUMMARY_WRITER_INIT_OP.setdefault(key, []).append(self._init_op)
def __init__(self, dataset): """Creates a new iterator over the given dataset. For example: ```python dataset = tf.data.Dataset.range(4) for x in Iterator(dataset): print(x) ``` Tensors produced will be placed on the device on which this iterator object was created. Args: dataset: A `tf.data.Dataset` object. Raises: RuntimeError: When invoked without eager execution enabled. """ if not context.executing_eagerly(): raise RuntimeError( "{} objects can only be used when eager execution is enabled, use " "tf.data.Dataset.make_initializable_iterator or " "tf.data.Dataset.make_one_shot_iterator for graph construction" .format(type(self))) with ops.device("/device:CPU:0"): ds_variant = dataset._as_variant_tensor() # pylint: disable=protected-access self._output_classes = dataset.output_classes self._output_types = dataset.output_types self._output_shapes = dataset.output_shapes self._flat_output_types = nest.flatten( sparse.as_dense_types(self._output_types, self._output_classes)) self._flat_output_shapes = nest.flatten( sparse.as_dense_shapes(self._output_shapes, self._output_classes)) self._resource = gen_dataset_ops.iterator( shared_name="", container=_generate_shared_name("eageriterator"), output_types=self._flat_output_types, output_shapes=self._flat_output_shapes) gen_dataset_ops.make_iterator(ds_variant, self._resource) # Delete the resource when this object is deleted self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device="/device:CPU:0") self._device = context.context().device_name
def __init__(self, dataset): """Creates a new iterator over the given dataset. For example: ```python dataset = tf.data.Dataset.range(4) for x in Iterator(dataset): print(x) ``` Tensors produced will be placed on the device on which this iterator object was created. Args: dataset: A `tf.data.Dataset` object. Raises: RuntimeError: When invoked without eager execution enabled. """ if not context.executing_eagerly(): raise RuntimeError( "{} objects can only be used when eager execution is enabled, use " "tf.data.Dataset.make_initializable_iterator or " "tf.data.Dataset.make_one_shot_iterator for graph construction" .format(type(self))) self._device = context.context().device_name with ops.device("/cpu:0"): # pylint: disable=protected-access dataset = dataset._apply_options() ds_variant = dataset._variant_tensor self._structure = structure_lib.convert_legacy_structure( dataset.output_types, dataset.output_shapes, dataset.output_classes) self._flat_output_types = self._structure._flat_types self._flat_output_shapes = self._structure._flat_shapes with ops.colocate_with(ds_variant): self._resource = gen_dataset_ops.anonymous_iterator( output_types=self._flat_output_types, output_shapes=self._flat_output_shapes) gen_dataset_ops.make_iterator(ds_variant, self._resource) # Delete the resource when this object is deleted self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device=self._device)
def __init__(self, shared_name, init_op_fn, name=None, v2=False, metadata=None): self._resource = gen_summary_ops.summary_writer( shared_name=shared_name, name=name) # TODO(nickfelt): cache other constructed ops in graph mode self._init_op_fn = init_op_fn self._init_op = init_op_fn(self._resource) self._v2 = v2 self._metadata = {} if metadata is None else metadata self._closed = False if context.executing_eagerly(): self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device="cpu:0") else: ops.add_to_collection(_SUMMARY_WRITER_INIT_COLLECTION_NAME, self._init_op)
def __init__(self, server_addresses): """Creates and starts the gRPC server. Args: server_addresses: A list of strings containing one or more server addresses. """ if not tf.executing_eagerly(): raise ValueError("Only eager mode is currently supported.") self._handle = gen_grpc_ops.grpc_server_resource_handle_op( shared_name=context.shared_name(None)) # Delete the resource when this object is deleted. self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._handle, handle_device=context.context().device_name) gen_grpc_ops.create_grpc_server(self._handle, server_addresses) # References to tf.Variable's, etc. used in a tf.function to prevent them # from being deallocated. self._keep_alive = []
def __init__(self, server_address): """Creates and starts the gRPC client. Args: server_address: A string containing the server address. """ if not tf.executing_eagerly(): raise ValueError("Only eager mode is currently supported.") self._handle = gen_grpc_ops.grpc_client_resource_handle_op( shared_name=context.shared_name(None)) self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._handle, handle_device=context.context().device_name) method_signatures = gen_grpc_ops.create_grpc_client(self._handle, server_address).numpy() m = service_pb2.MethodOutputSignature() v = struct_pb2.StructuredValue() for sig in method_signatures: assert m.ParseFromString(sig) decoder = nested_structure_coder.StructureCoder() assert v.ParseFromString(m.output_specs) decoded_output_specs = decoder.decode_proto(v) self._add_method(m.name, decoded_output_specs)
def __init__( self, # pylint: disable=super-init-not-called initial_value=None, trainable=None, caching_device=None, name=None, dtype=None, constraint=None, add_initializers_to=None, lifted_initializer_graph=None, synchronization=None, aggregation=None, **unused_kwargs): """Creates a variable. Args: initial_value: A `Tensor`, or Python object convertible to a `Tensor`, which is the initial value for the Variable. The initial value must have a shape specified unless `validate_shape` is set to False. Can also be a callable with no argument that returns the initial value when called. (Note that initializer functions from init_ops.py must first be bound to a shape before being used here.) trainable: If `True`, GradientTapes automatically watch uses of this Variable. caching_device: Optional device string or function describing where the Variable should be cached for reading. Defaults to the Variable's device. If not `None`, caches on another device. Typical use is to cache on the device where the Ops using the Variable reside, to deduplicate copying through `Switch` and other conditional statements. name: Optional name for the variable. Defaults to `'Variable'` and gets uniquified automatically. dtype: If set, initial_value will be converted to the given type. If None, either the datatype will be kept (if initial_value is a Tensor) or float32 will be used (if it is a Python object convertible to a Tensor). constraint: An optional projection function to be applied to the variable after being updated by an `Optimizer` (e.g. used to implement norm constraints or value constraints for layer weights). The function must take as input the unprojected Tensor representing the value of the variable and return the Tensor for the projected value (which must have the same shape). Constraints are not safe to use when doing asynchronous distributed training. add_initializers_to: if not None and not in legacy graph mode, the initializer tensor will be added to this map in addition to adding the assignment to the function. lifted_initializer_graph: FuncGraph to try to lift initializers to. synchronization: Indicates when a distributed a variable will be aggregated. Accepted values are constants defined in the class `tf.VariableSynchronization`. By default the synchronization is set to `AUTO` and the current `DistributionStrategy` chooses when to synchronize. If `synchronization` is set to `ON_READ`, `trainable` must not be set to `True`. aggregation: Indicates how a distributed variable will be aggregated. Accepted values are constants defined in the class `tf.VariableAggregation`. Raises: ValueError: If the initial value is not specified, or does not have a shape and `validate_shape` is `True`. RuntimeError: If called outside of a function definition. """ if not ops.inside_function(): # If we've been init_scope()d out of the function definition nothing to do # here; we can't really do the capturing or conditional logic. resource_variable_ops.ResourceVariable.__init__( self, initial_value=initial_value, trainable=trainable, caching_device=caching_device, name=name, dtype=dtype, constraint=constraint) return with ops.init_scope(): self._in_graph_mode = not context.executing_eagerly() if initial_value is None: raise ValueError("initial_value must be specified.") init_from_fn = callable(initial_value) if constraint is not None and not callable(constraint): raise ValueError("The `constraint` argument must be a callable.") if isinstance(initial_value, trackable.CheckpointInitialValue): self._maybe_initialize_trackable() self._update_uid = initial_value.checkpoint_position.restore_uid initial_value = initial_value.wrapped_value synchronization, aggregation, trainable = ( variables.validate_synchronization_aggregation_trainable( synchronization, aggregation, trainable, name)) self._trainable = trainable self._synchronization = synchronization self._aggregation = aggregation self._save_slice_info = None self._initial_value = None self._initializer_op = None self._is_initialized_op = None self._graph_element = None self._cached_value = None # Store the graph key so optimizers know how to only retrieve variables from # this graph. Guaranteed to be the same as the eager graph_key. self._graph_key = ops.get_default_graph()._graph_key # pylint: disable=protected-access with ops.name_scope(name, "Variable", [] if init_from_fn else [initial_value]) as name: # pylint: disable=protected-access with ops.init_scope(): handle_name = ops.name_from_scope_name(name) unique_id = "%s_%d" % (handle_name, ops.uid()) shared_name = context.shared_name(unique_id) with ops.name_scope("Initializer"), ops.device(None): initial_value = ops.convert_to_tensor( initial_value() if init_from_fn else initial_value, name="initial_value", dtype=dtype) with ops.init_scope(): self._handle = resource_variable_ops.eager_safe_variable_handle( initial_value=initial_value, shared_name=shared_name, name=name, graph_mode=self._in_graph_mode) self._shape = initial_value.shape self._unique_id = unique_id self._handle_name = handle_name + ":0" self._dtype = initial_value.dtype.base_dtype self._constraint = constraint assert initial_value is not None if self._in_graph_mode: with ops.init_scope(): outer_graph = ops.get_default_graph() func_graph = ops.get_default_graph() function_placeholders = (func_graph.inputs + func_graph.internal_captures) placeholder_ops = set( [tensor.op for tensor in function_placeholders]) lifted_initializer = lift_to_graph.lift_to_graph( [initial_value], outer_graph, disallowed_placeholders=placeholder_ops)[initial_value] with ops.init_scope(): self._initial_value = lifted_initializer with ops.name_scope("IsInitialized"): self._is_initialized_op = ( resource_variable_ops.var_is_initialized_op( self._handle)) if initial_value is not None: with ops.name_scope("Assign") as n, ops.colocate_with( self._handle): self._initializer_op = resource_variable_ops.assign_variable_op( self._handle, lifted_initializer, name=n) with ops.name_scope("Read"), ops.colocate_with( self._handle): # Manually assign reads to the handle's device to avoid log # messages. with ops.device(self._handle.device): value = self._read_variable_op() self._graph_element = value ops.add_to_collection(ops.GraphKeys.GLOBAL_VARIABLES, self) else: if add_initializers_to is not None: add_initializers_to[self] = initial_value def assign_fn(): with ops.name_scope("Assign") as n, ops.colocate_with( self._handle): resource_variable_ops.assign_variable_op(self._handle, initial_value, name=n) # Returning values to keep tf.cond happy. return ops.convert_to_tensor(1) def not_assign_fn(): return ops.convert_to_tensor(0) # Note: this cond is always guaranteed to run because we're inside a # defun which will insert automatic control dependencies. control_flow_ops.cond( resource_variable_ops.var_is_initialized_op(self._handle), not_assign_fn, assign_fn) # After the handle has been created, set up a way to clean it up when # executing eagerly. We'll hold the only reference to the deleter, so that # when this object is garbage collected the deleter will be too. This # means ResourceVariables can be part of reference cycles without those # cycles being uncollectable. if not self._in_graph_mode: self._handle_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._handle, handle_device=self._handle.device) self._cached_shape_as_list = None
def __init__(self, dataset, devices, max_buffer_size=1, prefetch_buffer_size=1, source_device="/cpu:0"): """Constructs a MultiDeviceIterator. Args: dataset: The input dataset to be iterated over. devices: The list of devices to fetch data to. max_buffer_size: Maximum size of the host side per device buffer to keep. prefetch_buffer_size: if > 1, then we setup a buffer on each device to prefetch into. source_device: The host device to place the `dataset` on. In order to prevent deadlocks, if the prefetch_buffer_size is greater than the max_buffer_size, we set the max_buffer_size to prefetch_buffer_size. """ options = dataset_ops.Options() options.experimental_distribute.num_devices = len(devices) dataset = dataset.with_options(options) self._dataset = dataset._apply_options() # pylint: disable=protected-access self._experimental_slack = dataset.options().experimental_slack self._devices = devices self._source_device = source_device self._source_device_tensor = ops.convert_to_tensor(source_device) self._max_buffer_size = max_buffer_size self._prefetch_buffer_size = prefetch_buffer_size if self._prefetch_buffer_size > self._max_buffer_size: self._max_buffer_size = self._prefetch_buffer_size # Create the MultiDeviceIterator. with ops.device(self._source_device): # TODO(b/121378567): Get rid of this shared_name hack. shared_name = "" if context.executing_eagerly(): shared_name = context.shared_name() self._multi_device_iterator_resource = ( gen_dataset_ops.multi_device_iterator( devices=self._devices, shared_name=shared_name, container="", **self._dataset._flat_structure)) # pylint: disable=protected-access if context.executing_eagerly(): # Delete the resource when this object is deleted self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._multi_device_iterator_resource, handle_device=self._source_device) # The incarnation ID is used to ensure consistency between the per-device # iterators and the multi-device iterator. self._incarnation_id = gen_dataset_ops.multi_device_iterator_init( self._dataset._variant_tensor, # pylint: disable=protected-access self._multi_device_iterator_resource, max_buffer_size=self._max_buffer_size) self._prototype_device_datasets = [] for i, device in enumerate(self._devices): with ops.device(device): ds = _PerDeviceGenerator(i, self._multi_device_iterator_resource, self._incarnation_id, self._source_device_tensor, self._dataset.element_spec) self._prototype_device_datasets.append(ds) # TODO(rohanj): Explore the possibility of the MultiDeviceIterator to # initialize the device side of the pipeline. This would allow the # MultiDeviceIterator to choose, for example, to move some transformations # into the device side from its input. It might be useful in rewriting. # Create the per device iterators. self._device_iterators = [] for i, device in enumerate(self._devices): with ops.device(device): ds = _create_device_dataset(self._prototype_device_datasets[i], self._incarnation_id, self._prefetch_buffer_size, self._experimental_slack) if context.executing_eagerly(): self._device_iterators.append( dataset_ops.make_one_shot_iterator(ds)) else: self._device_iterators.append( dataset_ops.make_initializable_iterator(ds)) if not context.executing_eagerly(): device_iterator_initializers = [ iterator.initializer for iterator in self._device_iterators ] self._initializer = control_flow_ops.group( *device_iterator_initializers)
def _set_up_resource_deleter(self): self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device="cpu:0")
def __init__(self, dataset, devices, max_buffer_size=1, prefetch_buffer_size=1, source_device="/cpu:0"): """Constructs a MultiDeviceIterator. Args: dataset: The input dataset to be iterated over. devices: The list of devices to fetch data to. max_buffer_size: Maximum size of the host side per device buffer to keep. prefetch_buffer_size: if > 1, then we setup a buffer on each device to prefetch into. source_device: The host device to place the `dataset` on. Raises: RuntimeError: If run in Eager mode. """ self._dataset = dataset._apply_options() # pylint: disable=protected-access self._devices = devices self._source_device = source_device self._source_device_tensor = ops.convert_to_tensor(source_device) # Create the MultiDeviceIterator. with ops.device(self._source_device): # TODO(b/121378567): Get rid of this shared_name hack. shared_name = "" if context.executing_eagerly(): shared_name = context.shared_name() self._multi_device_iterator_resource = ( gen_dataset_ops.multi_device_iterator( devices=self._devices, shared_name=shared_name, container="", **dataset_ops.flat_structure(dataset))) if context.executing_eagerly(): # Delete the resource when this object is deleted self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._multi_device_iterator_resource, handle_device=self._source_device) # The incarnation ID is used to ensure consistency between the per-device # iterators and the multi-device iterator. self._incarnation_id = gen_dataset_ops.multi_device_iterator_init( self._dataset._variant_tensor, # pylint: disable=protected-access self._multi_device_iterator_resource, max_buffer_size=max_buffer_size) # TODO(rohanj): Explore the possibility of the MultiDeviceIterator to # initialize the device side of the pipeline. This would allow the # MultiDeviceIterator to choose, for example, to move some transformations # into the device side from its input. It might be useful in rewriting. # Create the per device iterators. self._device_iterators = [] for i, device in enumerate(self._devices): with ops.device(device): ds = _PerDeviceGenerator(i, self._multi_device_iterator_resource, self._incarnation_id, self._source_device_tensor, dataset._element_structure) # pylint: disable=protected-access if prefetch_buffer_size > 0: ds = ds.prefetch(prefetch_buffer_size) # TODO(jsimsa): Enable auto-tuning and optimizations when supported for # non-CPU devices. options = dataset_ops.Options() options.experimental_autotune = False options.experimental_optimization.apply_default_optimizations = False ds = ds.with_options(options) if context.executing_eagerly(): self._device_iterators.append( dataset_ops.make_one_shot_iterator(ds)) else: self._device_iterators.append( dataset_ops.make_initializable_iterator(ds)) if not context.executing_eagerly(): device_iterator_initializers = [ iterator.initializer for iterator in self._device_iterators ] self._initializer = control_flow_ops.group( *device_iterator_initializers)
def __init__(self, dataset): """Creates a new iterator over the given dataset. For example: ```python dataset = tf.data.Dataset.range(4) for x in Iterator(dataset): print(x) ``` Tensors produced will be placed on the device on which this iterator object was created. Args: dataset: A `tf.data.Dataset` object. Raises: RuntimeError: When invoked without eager execution enabled. """ if not context.in_eager_mode(): raise RuntimeError( "{} objects can only be used when eager execution is enabled, use " "tf.data.Dataset.make_iterator or " "tf.data.Dataset.make_one_shot_iterator for graph construction" .format(type(self))) with ops.device("/device:CPU:0"): ds_variant = dataset._as_variant_tensor() # pylint: disable=protected-access self._output_types = dataset.output_types self._output_shapes = dataset.output_shapes self._flat_output_types = nest.flatten(dataset.output_types) self._flat_output_shapes = nest.flatten(dataset.output_shapes) self._resource = gen_dataset_ops.iterator( container="", shared_name=_generate_shared_name("eager_iterator"), output_types=self._flat_output_types, output_shapes=self._flat_output_shapes) gen_dataset_ops.make_iterator(ds_variant, self._resource) # Delete the resource when this object is deleted self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device="/device:CPU:0") self._device = context.context().device_name self._buffer_resource_handle = None if not context.context().device_spec.device_type: is_remote_device = False else: is_remote_device = context.context( ).device_spec.device_type != "CPU" if is_remote_device: with ops.device("/device:CPU:0"): iter_string_handle = gen_dataset_ops.iterator_to_string_handle( self._resource) @function.Defun(dtypes.string) def remote_fn(h): remote_iterator = iterator_ops.Iterator.from_string_handle( h, self._output_types, self._output_shapes) return remote_iterator.get_next() remote_fn.add_to_graph(None) target = constant_op.constant("/device:CPU:0") with ops.device(self._device): self._buffer_resource_handle = prefetching_ops.function_buffering_resource( string_arg=iter_string_handle, f=remote_fn, target_device=target, buffer_size=10, thread_pool_size=1, container="", shared_name=_generate_shared_name( "function_buffer_resource")) self._buffer_resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._buffer_resource_handle, handle_device=self._device)
def __init__(self, resource): self._resource = resource if context.in_eager_mode(): self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device="cpu:0")
def __init__(self, logdir, max_queue=10, flush_secs=120, filename_suffix=""): """Summary writer for TensorBoard, compatible with eager execution. If necessary, multiple instances of `SummaryWriter` can be created, with distinct `logdir`s and `name`s. Each `SummaryWriter` instance will retain its independent `global_step` counter and data writing destination. Example: ```python writer = tfe.SummaryWriter("my_model") # ... Code that sets up the model and data batches ... for _ in xrange(train_iters): loss = model.train_batch(batch) writer.scalar("loss", loss) writer.step() ``` Args: logdir: Directory in which summary files will be written. max_queue: Number of summary items to buffer before flushing to filesystem. If 0, summaries will be flushed immediately. flush_secs: Number of secondsbetween forced commits to disk. filename_suffix: Suffix of the event protobuf files in which the summary data are stored. Raises: ValueError: If this constructor is called not under eager execution. """ # TODO(apassos, ashankar): Make this class and the underlying # contrib.summary_ops compatible with graph model and remove this check. if not context.in_eager_mode(): raise ValueError( "Use of SummaryWriter is currently supported only with eager " "execution enabled. File an issue at " "https://github.com/tensorflow/tensorflow/issues/new to express " "interest in fixing this.") # TODO(cais): Consider adding name keyword argument, which if None or empty, # will register the global global_step that training_util.get_global_step() # can find. with context.device(self._CPU_DEVICE): self._name = uuid.uuid4().hex self._global_step = 0 self._global_step_tensor = variable_scope.get_variable( "global_step/summary_writer/" + self._name, shape=[], dtype=dtypes.int64, initializer=init_ops.zeros_initializer()) self._global_step_dirty = False self._resource = gen_summary_ops.summary_writer( shared_name=self._name) gen_summary_ops.create_summary_file_writer(self._resource, logdir, max_queue, flush_secs, filename_suffix) # Delete the resource when this object is deleted self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device=self._CPU_DEVICE)
def __init__(self, # pylint: disable=super-init-not-called initial_value=None, trainable=True, caching_device=None, name=None, dtype=None, constraint=None, **unused_kwargs): """Creates a variable. Args: initial_value: A `Tensor`, or Python object convertible to a `Tensor`, which is the initial value for the Variable. The initial value must have a shape specified unless `validate_shape` is set to False. Can also be a callable with no argument that returns the initial value when called. (Note that initializer functions from init_ops.py must first be bound to a shape before being used here.) trainable: If `True`, GradientTapes automatically watch uses of this Variable. caching_device: Optional device string or function describing where the Variable should be cached for reading. Defaults to the Variable's device. If not `None`, caches on another device. Typical use is to cache on the device where the Ops using the Variable reside, to deduplicate copying through `Switch` and other conditional statements. name: Optional name for the variable. Defaults to `'Variable'` and gets uniquified automatically. dtype: If set, initial_value will be converted to the given type. If None, either the datatype will be kept (if initial_value is a Tensor) or float32 will be used (if it is a Python object convertible to a Tensor). constraint: An optional projection function to be applied to the variable after being updated by an `Optimizer` (e.g. used to implement norm constraints or value constraints for layer weights). The function must take as input the unprojected Tensor representing the value of the variable and return the Tensor for the projected value (which must have the same shape). Constraints are not safe to use when doing asynchronous distributed training. Raises: ValueError: If the initial value is not specified, or does not have a shape and `validate_shape` is `True`. RuntimeError: If called outside of a function definition. """ if context.executing_eagerly(): raise RuntimeError( "UnliftedInitializerVariable should not be created " "outside of functions.") with ops.init_scope(): if not context.executing_eagerly(): raise RuntimeError( "UnliftedInitializerVariable does not support legacy graph mode.") self._in_graph_mode = False if initial_value is None: raise ValueError("initial_value must be specified.") init_from_fn = callable(initial_value) if constraint is not None and not callable(constraint): raise ValueError("The `constraint` argument must be a callable.") if isinstance(initial_value, checkpointable.CheckpointInitialValue): self._maybe_initialize_checkpointable() self._update_uid = initial_value.checkpoint_position.restore_uid initial_value = initial_value.wrapped_value self._trainable = trainable self._save_slice_info = None self._initial_value = None self._initializer_op = None self._is_initialized_op = None self._graph_element = None self._cached_value = None # Store the graph key so optimizers know how to only retrieve variables from # this graph. Guaranteed to be the same as the eager graph_key. self._graph_key = ops.get_default_graph()._graph_key # pylint: disable=protected-access with ops.name_scope(name, "Variable", [] if init_from_fn else [initial_value]) as name: # pylint: disable=protected-access with ops.init_scope(): assert context.executing_eagerly() shared_name = ops._name_from_scope_name(name) shared_name = "%s_%d" % (shared_name, ops.uid()) # Use attr_scope and device(None) to simulate the behavior of # colocate_with when the variable we want to colocate with doesn't # yet exist. with ops.name_scope("Initializer"), ops.device(None): initial_value = ops.convert_to_tensor( initial_value() if init_from_fn else initial_value, name="initial_value", dtype=dtype) with ops.init_scope(): self._handle = resource_variable_ops.eager_safe_variable_handle( shape=initial_value.get_shape(), dtype=initial_value.dtype.base_dtype, shared_name=shared_name, name=name, graph_mode=False) self._shape = initial_value.shape self._unique_id = shared_name self._handle_name = shared_name + ":0" self._dtype = initial_value.dtype.base_dtype self._constraint = constraint assert initial_value is not None def assign_fn(): with ops.name_scope("Assign") as n, ops.colocate_with(self._handle): resource_variable_ops.assign_variable_op( self._handle, initial_value, name=n) # Returning values to keep tf.cond happy. return ops.convert_to_tensor(1) def not_assign_fn(): return ops.convert_to_tensor(0) # Note: this cond is always guaranteed to run because we're inside a defun # which will insert automatic control dependencies. control_flow_ops.cond( resource_variable_ops.var_is_initialized_op(self._handle), not_assign_fn, assign_fn) # After the handle has been created, set up a way to clean it up when # executing eagerly. We'll hold the only reference to the deleter, so that # when this object is garbage collected the deleter will be too. This # means ResourceVariables can be part of reference cycles without those # cycles being uncollectable. self._handle_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._handle, handle_device=self._handle.device) self._cached_shape_as_list = None
def __init__(self, resource): self._resource = resource if context.executing_eagerly() and self._resource is not None: self._resource_deleter = resource_variable_ops.EagerResourceDeleter( handle=self._resource, handle_device="cpu:0")