| # Copyright 2018 The TensorFlow Authors. All Rights Reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # ============================================================================== |
| """Various classes representing distributed values.""" |
| |
| from __future__ import absolute_import |
| from __future__ import division |
| from __future__ import print_function |
| |
| import copy |
| import weakref |
| |
| from tensorflow.python.distribute import device_util |
| from tensorflow.python.distribute import distribute_lib |
| from tensorflow.python.distribute import distribution_strategy_context as ds_context |
| from tensorflow.python.distribute import packed_distributed_variable as packed |
| from tensorflow.python.distribute import reduce_util |
| from tensorflow.python.distribute import values_util |
| from tensorflow.python.eager import context |
| from tensorflow.python.framework import composite_tensor |
| from tensorflow.python.framework import ops |
| from tensorflow.python.framework import tensor_util |
| from tensorflow.python.framework import type_spec |
| from tensorflow.python.ops import array_ops |
| from tensorflow.python.ops import control_flow_ops |
| from tensorflow.python.ops import math_ops |
| from tensorflow.python.ops import variable_scope as vs |
| from tensorflow.python.ops import variables as variables_lib |
| from tensorflow.python.saved_model import save_context |
| from tensorflow.python.training.saving import saveable_object |
| from tensorflow.python.training.tracking import base as trackable |
| from tensorflow.python.types import core |
| from tensorflow.python.util.tf_export import tf_export |
| |
| |
| def _on_write_update_replica(var, update_fn, value, **kwargs): |
| """Updates variables with ON_WRITE synchronization in replica context.""" |
| if var.aggregation == vs.VariableAggregation.NONE: |
| return update_fn(var._get_on_device_or_primary(), value, **kwargs) # pylint: disable=protected-access |
| |
| if not getattr(ds_context.get_strategy().extended, "_use_merge_call", True): |
| # Don't allow MEAN with non float dtype, since it may cause unexpected |
| # precision loss. Python3 and NumPy automatically upcast integers to |
| # float in division, but we should always preserve the type. |
| if var.aggregation == vs.VariableAggregation.MEAN and ( |
| not var.dtype.is_floating) and tensor_util.is_tf_type(value): |
| raise ValueError( |
| "Cannot update non-float variables with " |
| "tf.VariableAggregation.MEAN aggregation in replica context. " |
| "Either change the variable dtype to float or update it in " |
| "cross-replica context.") |
| |
| aggregated_value = apply_aggregation_replica_context(value, var.aggregation, |
| var) |
| values_util.mark_as_unsaveable() |
| |
| return ds_context.get_replica_context()._update( # pylint: disable=protected-access |
| var, |
| update_fn, |
| args=(aggregated_value,), |
| kwargs=kwargs, group=True) |
| |
| else: |
| def merge_fn(strategy, value, **kwargs): |
| """Aggregate values and update all variables in cross replica context.""" |
| # Don't allow MEAN with non float dtype, since it may cause unexpected |
| # precision loss. Python3 and NumPy automatically upcast integers to |
| # float in division, but we should always preserve the type. |
| # |
| # Note that to be backward compatible we allow the case when the value |
| # is *always* the same on each replica. I.E. value is not a |
| # PerReplica. Refer to regroup() to see how values are grouped. |
| if var.aggregation == vs.VariableAggregation.MEAN and ( |
| not var.dtype.is_floating) and isinstance(value, PerReplica): |
| raise ValueError( |
| "Cannot update non-float variables with " |
| "tf.VariableAggregation.MEAN aggregation in replica context. " |
| "Either change the variable dtype to float or update it in " |
| "cross-replica context.") |
| |
| assert strategy == var.distribute_strategy |
| v = values_util.apply_aggregation(strategy, value, var.aggregation, var) |
| return var._update_cross_replica(update_fn, v, **kwargs) # pylint: disable=protected-access |
| |
| return ds_context.get_replica_context().merge_call( |
| merge_fn, args=(value,), kwargs=kwargs) |
| |
| |
| def apply_aggregation_replica_context(value, aggregation, destinations): |
| """Aggregate `value` to `destinations` as specified by `aggregation`.""" |
| # if it is a python literal, return without aggregation |
| if isinstance(value, DistributedValues): |
| raise ValueError("Do not use a DistributedValues update variable in replica" |
| " context.") |
| if not tensor_util.is_tf_type(value): |
| return value |
| |
| if aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA: |
| # Switch to cross-replica context to broadcast |
| def merge_fn(strategy, value): |
| return strategy.extended.broadcast_to( |
| strategy.experimental_local_results(value)[0], |
| destinations=destinations) |
| |
| return ds_context.get_replica_context().merge_call(merge_fn, args=(value,)) |
| |
| else: |
| reduce_op = reduce_util.ReduceOp.from_variable_aggregation(aggregation) |
| aggregated_value = ds_context.get_strategy( # pylint: disable=protected-access |
| ).extended._replica_ctx_all_reduce(reduce_op, value) |
| return aggregated_value |
| |
| |
| @tf_export("distribute.DistributedValues", v1=[]) |
| class DistributedValues(object): |
| """Base class for representing distributed values. |
| |
| A subclass instance of `tf.distribute.DistributedValues` is created when |
| creating variables within a distribution strategy, iterating a |
| `tf.distribute.DistributedDataset` or through `tf.distribute.Strategy.run`. |
| This base class should never be instantiated directly. |
| `tf.distribute.DistributedValues` contains a value per replica. Depending on |
| the subclass, the values could either be synced on update, synced on demand, |
| or never synced. |
| |
| `tf.distribute.DistributedValues` can be reduced to obtain single value across |
| replicas, as input into `tf.distribute.Strategy.run` or the per-replica values |
| inspected using `tf.distribute.Strategy.experimental_local_results`. |
| |
| Example usage: |
| |
| 1. Created from a `tf.distribute.DistributedDataset`: |
| |
| >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"]) |
| >>> dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2) |
| >>> dataset_iterator = iter(strategy.experimental_distribute_dataset(dataset)) |
| >>> distributed_values = next(dataset_iterator) |
| |
| 2. Returned by `run`: |
| |
| >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"]) |
| >>> @tf.function |
| ... def run(): |
| ... ctx = tf.distribute.get_replica_context() |
| ... return ctx.replica_id_in_sync_group |
| >>> distributed_values = strategy.run(run) |
| |
| 3. As input into `run`: |
| |
| >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"]) |
| >>> dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2) |
| >>> dataset_iterator = iter(strategy.experimental_distribute_dataset(dataset)) |
| >>> distributed_values = next(dataset_iterator) |
| >>> @tf.function |
| ... def run(input): |
| ... return input + 1.0 |
| >>> updated_value = strategy.run(run, args=(distributed_values,)) |
| |
| 4. Reduce value: |
| |
| >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"]) |
| >>> dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2) |
| >>> dataset_iterator = iter(strategy.experimental_distribute_dataset(dataset)) |
| >>> distributed_values = next(dataset_iterator) |
| >>> reduced_value = strategy.reduce(tf.distribute.ReduceOp.SUM, |
| ... distributed_values, |
| ... axis = 0) |
| |
| 5. Inspect local replica values: |
| |
| >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"]) |
| >>> dataset = tf.data.Dataset.from_tensor_slices([5., 6., 7., 8.]).batch(2) |
| >>> dataset_iterator = iter(strategy.experimental_distribute_dataset(dataset)) |
| >>> per_replica_values = strategy.experimental_local_results( |
| ... distributed_values) |
| >>> per_replica_values |
| (<tf.Tensor: shape=(1,), dtype=float32, numpy=array([5.], dtype=float32)>, |
| <tf.Tensor: shape=(1,), dtype=float32, numpy=array([6.], dtype=float32)>) |
| |
| """ |
| |
| def __init__(self, values): |
| """Should only be called by subclass __init__.""" |
| self._values = tuple(values) |
| |
| def _get(self): |
| """Returns the value for the current device or raises a ValueError.""" |
| replica_id = values_util.get_current_replica_id_as_int() |
| if replica_id is None: |
| return self._get_cross_replica() |
| else: |
| return self._values[replica_id] |
| |
| def _get_cross_replica(self): |
| raise NotImplementedError( |
| "This method should be overridden by sub-classes which support cross-" |
| "replica accesses.") |
| |
| def _get_on_device_or_primary(self): |
| """Returns value in same replica or device if possible, else the _primary.""" |
| replica_id = values_util.get_current_replica_id_as_int() |
| if replica_id is None: |
| # Try to find a value on the current device. |
| current_device = device_util.canonicalize(device_util.current()) |
| for value in self._values: |
| if device_util.canonicalize(value.device) == current_device: |
| return value |
| return self._primary |
| else: |
| return self._values[replica_id] |
| |
| @property |
| def _primary(self): |
| """Returns a representative component.""" |
| return self._values[0] |
| |
| @property |
| def _devices(self): |
| return tuple(v.device for v in self._values) |
| |
| def __str__(self): |
| debug_str = ",\n".join( |
| " %d: %s" % (i, v) for i, v in enumerate(self._values)) |
| return "%s:{\n%s\n}" % (self.__class__.__name__, debug_str) |
| |
| def __repr__(self): |
| debug_repr = ",\n".join( |
| " %d: %r" % (i, v) for i, v in enumerate(self._values)) |
| return "%s:{\n%s\n}" % (self.__class__.__name__, debug_repr) |
| |
| |
| # NOTE(josh11b,apassos): It would be great if we could inspect the values this was |
| # initialized with and use that to generate the overloaded operators here. |
| # Unfortunately, Python's rules for special methods don't allow this, see |
| # https://docs.python.org/3/reference/datamodel.html#special-method-names |
| # "if a class defines a method named __getitem__(), and x is an instance of |
| # this class, then x[i] is roughly equivalent to type(x).__getitem__(x, i)." |
| # In particular, these special methods don't go through __getattr__, and |
| # it will only use those methods if they are defined in the class, not the |
| # object. |
| class DistributedDelegate(DistributedValues): |
| """A map from device to values; acts as the same type as the values.""" |
| |
| def __getattr__(self, name): |
| # The '_use_resource_variables' and the attrs starts with '_self' are used |
| # for restoring the saved_model proto, and '_attribute_sentinel' is used for |
| # Layer tracking. At the point these attrs are queried, the variable has not |
| # been initialized. Thus it should not query those of the underlying |
| # components. |
| if name.startswith("_self_") or name in ("_use_resource_variables", |
| "_attribute_sentinel", |
| "_distributed_container"): |
| return super(DistributedDelegate, self).__getattr__(name) |
| |
| # This allows copy.copy(DistributedDelegate). When copying an object, |
| # copy.copy doesn't invoke its __init__ method, instead it makes a new |
| # empty object, then copies the attributes over. copy.copy looks for |
| # attributes like "__getstate__" in case the object implements its custom |
| # copying. Since DistributedDelegate doesn't have those attributes defined, |
| # __getattr__ will be invoked, which tries to access "_values" attributes, |
| # but that doesn't exist either because this is an empty object, and again |
| # __getattr__ is invoked, leading to an infinite recursion. |
| if name == "_values": |
| raise AttributeError() |
| |
| # TODO(priyag): This needs to be made robust against pitfalls from mix use |
| # __getattr__ and @property. See b/120402273. |
| return getattr(self._get(), name) |
| |
| @property |
| def values(self): |
| """Returns the per replica values.""" |
| return self._values |
| |
| def _get_as_operand(self): |
| """Returns the value for operations for the current device. |
| |
| Some implementations, e.g. `TPUMirroredVariable`, are not able to return the |
| value type within a replica context. They can, however, return a value that |
| can be used by the operations below. |
| """ |
| return self._get() |
| |
| # pylint: disable=multiple-statements |
| def __add__(self, o): |
| return self._get_as_operand() + o |
| |
| def __radd__(self, o): |
| return o + self._get_as_operand() |
| |
| def __sub__(self, o): |
| return self._get_as_operand() - o |
| |
| def __rsub__(self, o): |
| return o - self._get_as_operand() |
| |
| def __mul__(self, o): |
| return self._get_as_operand() * o |
| |
| def __rmul__(self, o): |
| return o * self._get_as_operand() |
| |
| def __truediv__(self, o): |
| return self._get_as_operand() / o |
| |
| def __rtruediv__(self, o): |
| return o / self._get_as_operand() |
| |
| def __floordiv__(self, o): |
| return self._get_as_operand() // o |
| |
| def __rfloordiv__(self, o): |
| return o // self._get_as_operand() |
| |
| def __mod__(self, o): |
| return self._get_as_operand() % o |
| |
| def __rmod__(self, o): |
| return o % self._get_as_operand() |
| |
| def __lt__(self, o): |
| return self._get_as_operand() < o |
| |
| def __le__(self, o): |
| return self._get_as_operand() <= o |
| |
| def __gt__(self, o): |
| return self._get_as_operand() > o |
| |
| def __ge__(self, o): |
| return self._get_as_operand() >= o |
| |
| def __and__(self, o): |
| return self._get_as_operand() & o |
| |
| def __rand__(self, o): |
| return o & self._get_as_operand() |
| |
| def __or__(self, o): |
| return self._get_as_operand() | o |
| |
| def __ror__(self, o): |
| return o | self._get_as_operand() |
| |
| def __xor__(self, o): |
| return self._get_as_operand() ^ o |
| |
| def __rxor__(self, o): |
| return o ^ self._get_as_operand() |
| |
| def __getitem__(self, o): |
| return self._get_as_operand()[o] |
| |
| def __pow__(self, o, modulo=None): |
| return pow(self._get_as_operand(), o, modulo) |
| |
| def __rpow__(self, o): |
| return pow(o, self._get_as_operand()) |
| |
| def __invert__(self): |
| return ~self._get_as_operand() |
| |
| def __neg__(self): |
| return -self._get_as_operand() |
| |
| def __abs__(self): |
| return abs(self._get_as_operand()) |
| |
| def __div__(self, o): |
| try: |
| return self._get_as_operand().__div__(o) |
| except AttributeError: |
| # See https://docs.python.org/3/library/constants.html#NotImplemented |
| return NotImplemented |
| |
| def __rdiv__(self, o): |
| try: |
| return self._get_as_operand().__rdiv__(o) |
| except AttributeError: |
| # See https://docs.python.org/3/library/constants.html#NotImplemented |
| return NotImplemented |
| |
| def __matmul__(self, o): |
| try: |
| return self._get_as_operand().__matmul__(o) |
| except AttributeError: |
| # See https://docs.python.org/3/library/constants.html#NotImplemented |
| return NotImplemented |
| |
| def __rmatmul__(self, o): |
| try: |
| return self._get_as_operand().__rmatmul__(o) |
| except AttributeError: |
| # See https://docs.python.org/3/library/constants.html#NotImplemented |
| return NotImplemented |
| |
| # TODO(josh11b): Even more operator overloads. |
| |
| |
| class PerReplica(DistributedValues, composite_tensor.CompositeTensor): |
| """Holds a map from replica to unsynchronized values.""" |
| |
| @property |
| def _type_spec(self): |
| return PerReplicaSpec( |
| *(type_spec.type_spec_from_value(v) for v in self._values)) |
| |
| @property |
| def values(self): |
| """Returns the per replica values.""" |
| return self._values |
| |
| |
| class PerReplicaSpec(type_spec.TypeSpec): |
| """Type specification for a `PerReplica`.""" |
| |
| __slots__ = ["_value_specs"] |
| |
| value_type = property(lambda self: PerReplica) |
| |
| def __init__(self, *value_specs): |
| self._value_specs = tuple(value_specs) |
| |
| def _serialize(self): |
| return self._value_specs |
| |
| @property |
| def _component_specs(self): |
| return self._value_specs |
| |
| def _to_components(self, value): |
| replica_context = ds_context.get_replica_context() |
| if replica_context is not None and replica_context.num_replicas_in_sync > 1: |
| raise ValueError( |
| "Flattening a PerReplica to components is not supported in replica " |
| "context.") |
| return value._values # pylint: disable=protected-access |
| |
| def _from_components(self, tensor_list): |
| return PerReplica(tensor_list) |
| |
| |
| # Note that unlike PerReplica, Mirrored values inherit from |
| # DistributedDelegate and so can be used directly in cross-replica mode. |
| # TODO(tomhennigan) Should this extend CompositeTensor? |
| class Mirrored(DistributedDelegate): |
| """Holds a map from replica to values which are kept in sync.""" |
| |
| def _get_cross_replica(self): |
| return self._get_on_device_or_primary() |
| |
| def _as_graph_element(self): |
| obj = self._get() |
| conv_fn = getattr(obj, "_as_graph_element", None) |
| if conv_fn and callable(conv_fn): |
| return conv_fn() |
| return obj |
| |
| |
| class DistributedVarOp(object): |
| """A class that looks like `tf.Operation`.""" |
| |
| def __init__(self, name, graph, traceback, typ): |
| self.name = name |
| self.graph = graph |
| self.traceback = traceback |
| self.type = typ |
| |
| def __eq__(self, o): |
| if not isinstance(o, self.__class__): |
| raise NotImplementedError |
| return (self.name == o.name and self.graph == o.graph and |
| self.traceback == o.traceback and self.type == o.type) |
| |
| def __hash__(self): |
| return hash((self.name, self.graph, tuple(self.traceback), self.type)) |
| |
| |
| class DistributedVariable(DistributedDelegate, variables_lib.Variable, |
| core.Tensor): |
| """Holds a map from replica to variables.""" |
| |
| def __init__(self, strategy, values, aggregation, var_policy=None): |
| if (aggregation == variables_lib.VariableAggregation.MEAN and |
| not values[0].dtype.is_floating): |
| raise ValueError( |
| "creating distributed tf.Variable with aggregation=MEAN and a " |
| "non-floating dtype is not supported, please use a different " |
| "aggregation or dtype") |
| self._distribute_strategy = strategy |
| self._aggregation = aggregation |
| super(DistributedVariable, self).__init__(values) |
| self._common_name = self._primary.name.split(":")[0] |
| # Use a weakref to make it easy to map from the contained values |
| # to the container without introducing a reference cycle. |
| for v in values: |
| v._distributed_container = weakref.ref(self) # pylint: disable=protected-access |
| |
| # Packed variable is used to reduce the overhead of function execution. |
| # For a DistributedVariable, only one variable handle is captured into a |
| # function graph. It's only supported in eager mode. |
| if ops.executing_eagerly_outside_functions() and getattr( |
| strategy, "_enable_packed_variable_in_eager_mode", False): |
| name = "%s/packed/" % self._common_name |
| self._packed_var = packed.PackedDistributedVariable(values, name=name) |
| else: |
| self._packed_var = None |
| |
| # tf.keras keeps track of variables initialized using this attribute. When |
| # tf.keras gets the default session, it initializes all uninitialized vars. |
| # We need to make _keras_initialized a member of DistributedVariable because |
| # without this it will use `__getattr__` which will delegate to a component |
| # variable. |
| self._keras_initialized = False |
| # Typically, a `DistributedVariable`'s initializer is composed of the |
| # initializers of the components variables. However, in some cases, such as |
| # when restoring from a checkpoint, we may set the _initializer_op |
| # property on the entire `DistributedVariable`. |
| self._initializer_op = None |
| # Set a VariablePolicy which decides how we replicate/aggregate the given |
| # variable. |
| self._policy = var_policy |
| |
| def __deepcopy__(self, memo): |
| """Perform a deepcopy of the `DistributedVariable`. |
| |
| Unlike the deepcopy of a regular tf.Variable, this keeps the original |
| strategy and devices of the `DistributedVariable`. To avoid confusion |
| with the behavior of deepcopy on a regular `Variable` (which does |
| copy into new devices), we only allow a deepcopy of a `DistributedVariable` |
| within its originating strategy scope. |
| |
| Args: |
| memo: The memoization object for `deepcopy`. |
| |
| Returns: |
| A deep copy of the current `DistributedVariable`. |
| |
| Raises: |
| RuntimeError: If trying to deepcopy into a different strategy. |
| """ |
| with ds_context.enter_or_assert_strategy(self._distribute_strategy): |
| new_values = [] |
| |
| for value in self._values: |
| with ops.device(value.device): |
| new_values.append(copy.deepcopy(value, memo)) |
| |
| copied_variable = type(self)( |
| strategy=self._distribute_strategy, |
| values=new_values, |
| aggregation=self._aggregation, |
| var_policy=copy.deepcopy(self._policy, memo)) |
| |
| memo[id(self)] = copied_variable |
| |
| return copied_variable |
| |
| def _use_packed_variable(self): |
| # Don't use packed variable when under a SaveContext to avoid explicit |
| # device placement on variable consuming ops. |
| return self._packed_var is not None and not save_context.in_save_context() |
| |
| def is_initialized(self, name=None): |
| """Identifies if all the component variables are initialized. |
| |
| Args: |
| name: Name of the final `logical_and` op. |
| |
| Returns: |
| The op that evaluates to True or False depending on if all the |
| component variables are initialized. |
| """ |
| if values_util.is_saving_non_distributed(): |
| return self._primary.is_initialized() |
| if self._use_packed_variable(): |
| return self._packed_var.is_initialized() |
| result = self._primary.is_initialized() |
| # We iterate through the list of values except the last one to allow us to |
| # name the final `logical_and` op the same name that is passed by the user |
| # to the `is_initialized` op. For distributed variables, the |
| # `is_initialized` op is a `logical_and` op. |
| for v in self._values[1:-1]: |
| result = math_ops.logical_and(result, v.is_initialized()) |
| result = math_ops.logical_and( |
| result, self._values[-1].is_initialized(), name=name) |
| return result |
| |
| @property |
| def initializer(self): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.initializer |
| if self._initializer_op: |
| init_op = self._initializer_op |
| else: |
| # return grouped ops of all the var initializations of component values of |
| # the mirrored variable |
| init_op = control_flow_ops.group( |
| tuple(v.initializer for v in self._values)) |
| return init_op |
| |
| def initialized_value(self): |
| return self._get_on_device_or_primary().initialized_value() |
| |
| @property |
| def initial_value(self): |
| return self._get_on_device_or_primary().initial_value |
| |
| @property |
| def constraint(self): |
| return self._primary.constraint |
| |
| @property |
| def graph(self): |
| return self._primary.graph |
| |
| @property |
| def _shared_name(self): |
| return self._common_name |
| |
| @property |
| def _unique_id(self): |
| return self._primary._unique_id # pylint: disable=protected-access |
| |
| @property |
| def _graph_key(self): |
| """Lets Optimizers know which graph this variable is from.""" |
| return self._primary._graph_key # pylint: disable=protected-access |
| |
| @property |
| def name(self): |
| return self._primary.name |
| |
| @property |
| def dtype(self): |
| return self._primary.dtype |
| |
| @property |
| def shape(self): |
| return self._primary.shape |
| |
| @property |
| def synchronization(self): |
| return self._primary.synchronization |
| |
| @property |
| def aggregation(self): |
| return self._aggregation |
| |
| @property |
| def _packed_variable(self): |
| if self._use_packed_variable(): |
| return self._packed_var |
| return None |
| |
| @property |
| def handle(self): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.handle |
| replica_id = values_util.get_current_replica_id_as_int() |
| if replica_id is None: |
| raise ValueError("`handle` is not available outside the replica context" |
| " or a `tf.distribute.Strategy.update()` call.") |
| else: |
| if self._use_packed_variable(): |
| return self._packed_var.handle |
| return self._values[replica_id].handle |
| |
| def eval(self, session=None): |
| return self._get_on_device_or_primary().eval(session) |
| |
| @property |
| def _save_slice_info(self): |
| return self._primary._save_slice_info # pylint: disable=protected-access |
| |
| def _get_save_slice_info(self): |
| return self._primary._get_save_slice_info() # pylint: disable=protected-access |
| |
| def _set_save_slice_info(self, save_slice_info): |
| for v in self._values: |
| v._set_save_slice_info(save_slice_info) # pylint: disable=protected-access |
| |
| @property |
| def device(self): |
| return self._get_on_device_or_primary().device |
| |
| @property |
| def trainable(self): |
| return self._primary.trainable |
| |
| @property |
| def distribute_strategy(self): |
| return self._distribute_strategy |
| |
| def get_shape(self): |
| return self._primary.get_shape() |
| |
| def to_proto(self, export_scope=None): |
| return self._primary.to_proto(export_scope=export_scope) |
| |
| @property |
| def op(self): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.op |
| # We want cross-replica code that does some var.op.X calls |
| # to work (even if the current device isn't in self._devices), but |
| # other uses of var.op in a cross-replica context to fail. |
| if ds_context.in_cross_replica_context(): |
| return DistributedVarOp(self._primary.op.name, self._primary.op.graph, |
| self._primary.op.traceback, self._primary.op.type) |
| return self._get().op |
| |
| @property |
| def _in_graph_mode(self): |
| return self._primary._in_graph_mode # pylint: disable=protected-access |
| |
| def _get_replica(self, replica_id): |
| """Returns the value on a device with the given replica_id.""" |
| if self._use_packed_variable(): |
| return self._packed_var.on_device(self._devices[replica_id]) |
| return self._values[replica_id] |
| |
| def _get(self): |
| """Returns the value for the current device or raises a ValueError.""" |
| if values_util.is_saving_non_distributed(): |
| return self._primary |
| replica_id = values_util.get_current_replica_id_as_int() |
| if replica_id is None: |
| return self._get_cross_replica() |
| else: |
| return self._get_replica(replica_id) |
| |
| def _get_on_device_or_primary(self): |
| """Returns value in same replica or device if possible, else the _primary.""" |
| if values_util.is_saving_non_distributed(): |
| return self._primary |
| replica_id = values_util.get_current_replica_id_as_int() |
| if replica_id is None: |
| # Try to find a value on the current device. |
| current_device = device_util.canonicalize(device_util.current()) |
| for i, value in enumerate(self._values): |
| if device_util.canonicalize(value.device) == current_device: |
| return self._get_replica(i) |
| return self._get_replica(0) |
| else: |
| return self._get_replica(replica_id) |
| |
| def read_value(self): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.read_value() |
| with ds_context.enter_or_assert_strategy(self._distribute_strategy): |
| return array_ops.identity(self._get()) |
| |
| def value(self): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.value() |
| if self._policy: |
| return self._policy.value(self) |
| return self._get_on_device_or_primary().value() |
| |
| def numpy(self): |
| if context.executing_eagerly(): |
| return self.read_value().numpy() |
| else: |
| raise NotImplementedError( |
| "numpy() is only available when eager execution is enabled.") |
| |
| def assign_sub(self, value, use_locking=False, name=None, read_value=True): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.assign_sub(value, use_locking, name, read_value) |
| if self._policy: |
| return self._policy.assign_sub( |
| self, |
| value, |
| use_locking=use_locking, |
| name=name, |
| read_value=read_value) |
| return values_util.on_write_assign_sub( |
| self, value, use_locking=use_locking, name=name, read_value=read_value) |
| |
| def assign_add(self, value, use_locking=False, name=None, read_value=True): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.assign_add(value, use_locking, name, read_value) |
| if self._policy: |
| return self._policy.assign_add( |
| self, |
| value, |
| use_locking=use_locking, |
| name=name, |
| read_value=read_value) |
| return values_util.on_write_assign_add( |
| self, value, use_locking=use_locking, name=name, read_value=read_value) |
| |
| def assign(self, value, use_locking=False, name=None, read_value=True): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.assign(value, use_locking, name, read_value) |
| if self._policy: |
| return self._policy.assign( |
| self, |
| value, |
| use_locking=use_locking, |
| name=name, |
| read_value=read_value) |
| return values_util.on_write_assign( |
| self, value, use_locking=use_locking, name=name, read_value=read_value) |
| |
| def scatter_sub(self, sparse_delta, use_locking=False, name=None): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_sub(sparse_delta, use_locking, name) |
| if self._policy: |
| return self._policy.scatter_sub( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| return values_util.scatter_sub( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| |
| def scatter_add(self, sparse_delta, use_locking=False, name=None): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_add(sparse_delta, use_locking, name) |
| if self._policy: |
| return self._policy.scatter_add( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| return values_util.scatter_add( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| |
| def scatter_mul(self, sparse_delta, use_locking=False, name=None): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_mul(sparse_delta, use_locking, name) |
| if self._policy: |
| return self._policy.scatter_mul( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| return values_util.scatter_mul( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| |
| def scatter_div(self, sparse_delta, use_locking=False, name=None): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_div(sparse_delta, use_locking, name) |
| if self._policy: |
| return self._policy.scatter_div( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| return values_util.scatter_div( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| |
| def scatter_min(self, sparse_delta, use_locking=False, name=None): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_min(sparse_delta, use_locking, name) |
| if self._policy: |
| return self._policy.scatter_min( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| return values_util.scatter_min( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| |
| def scatter_max(self, sparse_delta, use_locking=False, name=None): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_max(sparse_delta, use_locking, name) |
| if self._policy: |
| return self._policy.scatter_max( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| return values_util.scatter_max( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| |
| def scatter_update(self, sparse_delta, use_locking=False, name=None): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_update(sparse_delta, use_locking, name) |
| if self._policy: |
| return self._policy.scatter_update( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| return values_util.scatter_update( |
| self, sparse_delta, use_locking=use_locking, name=name) |
| |
| def _gather_saveables_for_checkpoint(self): |
| """Overrides Trackable method. |
| |
| This allows both name-based and object-based save and restore of |
| DistributedVariables. |
| |
| Returns: |
| A dictionary mapping attribute names to `SaveableObject` factories. |
| """ |
| |
| def _saveable_factory(name=self._common_name): |
| return _DistributedVariableSaveable(self, self._primary, name) |
| |
| return {trackable.VARIABLE_VALUE_KEY: _saveable_factory} |
| |
| def _as_graph_element(self): |
| if values_util.is_saving_non_distributed(): |
| return self._primary._as_graph_element() # pylint: disable=protected-access |
| if self._policy: |
| return self._policy._as_graph_element(self) # pylint: disable=protected-access |
| |
| raise NotImplementedError("No policy set for calling _as_graph_element.") |
| |
| def _get_cross_replica(self): |
| if values_util.is_saving_non_distributed(): |
| return self._primary |
| if self._policy: |
| return self._policy._get_cross_replica(self) # pylint: disable=protected-access |
| |
| raise NotImplementedError( |
| "This method should be overridden by sub-classes which support cross-" |
| "replica accesses.") |
| |
| def _update_cross_replica(self, update_fn, value, **kwargs): |
| """Applies updates across replicas. |
| |
| Args: |
| update_fn: A callable to pass to `strategy.extended.update` to update the |
| variable. It should has the same signature as `Variable.assign()`. |
| value: value to be passed to `update_fn`. |
| **kwargs: remaining arguments to `update_fn`. |
| |
| Returns: |
| Updated variable or `tf.Operation`. |
| """ |
| values_util.mark_as_unsaveable() |
| return self.distribute_strategy.extended.update( |
| self, update_fn, args=(value,), kwargs=kwargs, group=True) |
| |
| def _update_replica(self, update_fn, value, **kwargs): |
| """Applies updates in one replica. |
| |
| Args: |
| update_fn: A callable to update the variable. It should has the same |
| signature as `Variable.assign()`. |
| value: value to be passed to `update_fn`. |
| **kwargs: remaining arguments to `update_fn`. |
| |
| Returns: |
| Updated variable or `tf.Operation`. |
| """ |
| if self._policy: |
| return self._policy._update_replica(self, update_fn, value, **kwargs) # pylint: disable=protected-access |
| raise NotImplementedError("should be implemented by subclass.") |
| |
| def _update(self, update_fn, value, **kwargs): |
| """Applies updates depending on the context. |
| |
| The method calls `_update_replica` in replica context, |
| `_update_cross_replica` in cross replica context, and `update_fn` in update |
| context. |
| |
| If `read_value` is True, the method returns the updated Variable. If |
| `read_value` is False, the method returns the update `tf.Operation`. |
| |
| Args: |
| update_fn: A callable to pass to `strategy.extended.update` to update the |
| variable. It should have the same signature as `Variable.assign()`. |
| value: value to be passed to `update_fn`. |
| **kwargs: keyword arguments to `update_fn`. |
| |
| Returns: |
| Updated variable or `tf.Operation`. |
| |
| """ |
| if values_util.is_saving_non_distributed(): |
| return update_fn(self._primary, value, **kwargs) |
| with ds_context.enter_or_assert_strategy(self.distribute_strategy): |
| if ds_context.in_cross_replica_context(): |
| update_replica_id = distribute_lib.get_update_replica_id() |
| if update_replica_id is not None: |
| replica_value = self._get_replica(update_replica_id) |
| return update_fn(replica_value, value, **kwargs) |
| return self._update_cross_replica(update_fn, value, **kwargs) |
| else: |
| values_util.assert_replica_context(self.distribute_strategy) |
| return self._update_replica(update_fn, value, **kwargs) |
| |
| def _should_act_as_resource_variable(self): |
| """Pass resource_variable_ops.is_resource_variable check.""" |
| pass |
| |
| def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False): |
| """Converts a variable to a tensor.""" |
| if values_util.is_saving_non_distributed(): |
| return ops.convert_to_tensor( |
| self._primary, dtype=dtype, name=name, as_ref=as_ref) |
| with ds_context.enter_or_assert_strategy(self._distribute_strategy): |
| return ops.convert_to_tensor( |
| self._get(), dtype=dtype, name=name, as_ref=as_ref) |
| |
| def _map_resources(self, save_options): |
| """For implementing `Trackable`.""" |
| # Initialize for self._primary first, so that obj_map[self._primary] and |
| # resource_map[self._primary.handle] contain mapped values. |
| obj_map, resource_map = self._primary._map_resources(save_options) # pylint:disable=protected-access |
| for v in [v for v in self._values if v != self._primary]: |
| |
| if (save_options.experimental_variable_policy # pylint:disable=protected-access |
| ._expand_distributed_variables()): |
| v_obj_map, v_resource_map = v._map_resources(save_options) # pylint:disable=protected-access |
| obj_map.update(v_obj_map) |
| resource_map.update(v_resource_map) |
| else: |
| obj_map[v] = obj_map[self._primary] |
| resource_map[v.handle] = resource_map[self._primary.handle] |
| obj_map[self] = obj_map[self._primary] |
| resource_map[self] = resource_map[self._primary.handle] |
| if self._packed_var is not None: |
| resource_map[self._packed_var.packed_handle] = resource_map[ |
| self._primary.handle] |
| return obj_map, resource_map |
| |
| def _write_object_proto(self, proto, options): |
| """Update a SavedObject proto for the caller. |
| |
| If a DistributedVariable object supports this method, it will be called when |
| saving with a pre-built `SavedObject` proto representing the object, plus an |
| instance of `SaveOptions`. This method is then free to modify that proto |
| instance. |
| |
| `DistributedVariable` with `AUTO` or `ON_WRITE` synchronization optionally |
| write out information about their components to the |
| `experimental_distributed_variable_components` field of a |
| `SavedVariable` (depending on the `SaveOptions` variable policy). |
| |
| Args: |
| proto: A pre-built `SavedObject` proto for this object. It is assumed this |
| will be a `SavedVariable` instance. |
| options: A `SaveOptions` instance. |
| """ |
| if self._policy: |
| if self._policy._is_mirrored(): # pylint: disable=protected-access |
| self._policy._write_object_proto(self, proto, options) # pylint: disable=protected-access |
| else: |
| self._write_object_proto(proto, options) |
| |
| |
| # We extend from `saveable_object.SaveableObject` instead of |
| # `saveable_object_util.ResourceVariableSaveable` since we need to read the |
| # value of ONREAD variables when saving. `SaveableObject` provides a way to |
| # specify the function to run to get the value of the variable or tensor at |
| # saving time. We can use this for both ON_READ and ON_WRITE variables. |
| # TODO(b/164586507): Consolidate ON_WRITE and ON_READ saving/restoring logic |
| # if possible. |
| class _DistributedVariableSaveable(saveable_object.SaveableObject): |
| """Class for defining how to restore a DistributedVariable.""" |
| |
| def __init__(self, distributed_variable, primary_variable, name): |
| self._distributed_variable = distributed_variable |
| if not self._distributed_variable._policy: |
| raise ValueError("VariablePolicy has not been set for the distributed " |
| "variable.") |
| tensor, spec = distributed_variable._policy.get_saveable( |
| distributed_variable, primary_variable, name) |
| super(_DistributedVariableSaveable, self).__init__(tensor, spec, name) |
| |
| def restore(self, restored_tensors, restored_shapes): |
| """Restore the same value into all variables.""" |
| tensor, = restored_tensors |
| return self._distributed_variable._policy.get_restore_ops( # pylint: disable=protected-access |
| self._distributed_variable, tensor) |
| |
| |
| class _MirroredSaveable(saveable_object.SaveableObject): |
| """Class for defining how to restore a MirroredVariable.""" |
| |
| def __init__(self, mirrored_variable, primary_variable, name): |
| self._mirrored_variable = mirrored_variable |
| tensor, spec = values_util.get_on_write_saveable(self._mirrored_variable, |
| primary_variable, |
| name) |
| super(_MirroredSaveable, self).__init__(tensor, spec, name) |
| |
| def restore(self, restored_tensors, restored_shapes): |
| """Restore the same value into all variables.""" |
| tensor, = restored_tensors |
| return values_util.get_on_write_restore_ops(self._mirrored_variable, |
| tensor) |
| |
| |
| class MirroredVariable(DistributedVariable, Mirrored): |
| """Holds a map from replica to variables whose values are kept in sync.""" |
| |
| def _update_replica(self, update_fn, value, **kwargs): |
| return _on_write_update_replica(self, update_fn, value, **kwargs) |
| |
| def scatter_min(self, *args, **kwargs): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_min(*args, **kwargs) |
| if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and |
| self._aggregation != vs.VariableAggregation.NONE): |
| raise NotImplementedError(values_util.scatter_error_msg.format( |
| op_name="scatter_min", aggregation=self._aggregation)) |
| return super(MirroredVariable, self).scatter_min(*args, **kwargs) |
| |
| def scatter_max(self, *args, **kwargs): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_max(*args, **kwargs) |
| if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and |
| self._aggregation != vs.VariableAggregation.NONE): |
| raise NotImplementedError(values_util.scatter_error_msg.format( |
| op_name="scatter_max", aggregation=self._aggregation)) |
| return super(MirroredVariable, self).scatter_max(*args, **kwargs) |
| |
| def scatter_update(self, *args, **kwargs): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_update(*args, **kwargs) |
| if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and |
| self._aggregation != vs.VariableAggregation.NONE): |
| raise NotImplementedError(values_util.scatter_error_msg.format( |
| op_name="scatter_update", aggregation=self._aggregation)) |
| return super(MirroredVariable, self).scatter_update(*args, **kwargs) |
| |
| def _get_cross_replica(self): |
| # Return identity, to avoid directly exposing the variable to the user and |
| # allowing it to be modified by mistake. |
| return array_ops.identity(Mirrored._get_cross_replica(self)) |
| |
| def _as_graph_element(self): |
| return self._get_on_device_or_primary()._as_graph_element() # pylint: disable=protected-access |
| |
| def _gather_saveables_for_checkpoint(self): |
| """Overrides Trackable method. |
| |
| This allows both name-based and object-based save and restore of |
| MirroredVariables. |
| |
| Returns: |
| A dictionary mapping attribute names to `SaveableObject` factories. |
| """ |
| |
| def _saveable_factory(name=self._common_name): |
| return _MirroredSaveable(self, self._primary, name) |
| |
| return {trackable.VARIABLE_VALUE_KEY: _saveable_factory} |
| |
| def _write_object_proto(self, proto, options): |
| """Update a SavedObject proto for the caller. |
| |
| If a DistributedVariable object supports this method, it will be called when |
| saving with a pre-built `SavedObject` proto representing the object, plus an |
| instance of `SaveOptions`. This method is then free to modify that proto |
| instance. |
| |
| `DistributedVariable` with `AUTO` or `ON_WRITE` synchronization optionally |
| write out information about their components to the |
| `experimental_distributed_variable_components` field of a |
| `SavedVariable` (depending on the `SaveOptions` variable policy). |
| |
| Args: |
| proto: A pre-built `SavedObject` proto for this object. It is assumed this |
| will be a `SavedVariable` instance. |
| options: A `SaveOptions` instance. |
| """ |
| values_util.write_object_proto(self, proto, options) |
| |
| def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False): |
| """Converts a variable to a tensor.""" |
| # TODO(b/154017756): Make _dense_var_to_tensor consistent between ON_READ |
| # and ON_WRITE. |
| # Try to avoid assignments to and other mutations of MirroredVariable |
| # state except through a DistributionStrategy.extended.update() or any of |
| # the `assign*` and `scatter*` calls. |
| if as_ref: |
| # A TF 1.x case where the variable is a boolean variable and used like: |
| # tf.cond(v, true_fn, false_fn). |
| raise ValueError( |
| "You may be using variable created under distribute strategy in TF " |
| "1.x control flows. Try explicitly converting the variable to Tensor " |
| "using variable.read_value(), or switch to TF 2.x.") |
| return ops.convert_to_tensor( |
| self._get(), dtype=dtype, name=name, as_ref=as_ref) |
| |
| |
| class _SyncOnReadSaveable(saveable_object.SaveableObject): |
| """Class for defining how to restore a SyncOnReadVariable.""" |
| |
| def __init__(self, sync_on_read_variable, name): |
| self._sync_on_read_variable = sync_on_read_variable |
| tensor, spec = values_util.get_on_read_saveable( |
| sync_on_read_variable, sync_on_read_variable._primary, name) |
| |
| super(_SyncOnReadSaveable, self).__init__(tensor, spec, name) |
| |
| def restore(self, restored_tensors, restored_shapes): |
| """Restore the same value into all variables.""" |
| tensor, = restored_tensors |
| return values_util.get_on_read_restore_ops( |
| self._sync_on_read_variable, tensor, |
| self._sync_on_read_variable.aggregation) |
| |
| |
| class SyncOnReadVariable(DistributedVariable): |
| """Holds a map from replica to variables whose values are reduced on save.""" |
| |
| def _update_replica(self, update_fn, value, **kwargs): |
| return update_fn(self._get_on_device_or_primary(), value, **kwargs) |
| |
| # TODO(b/154017756): Make assign behaivor in cross replica context consistent |
| # with MirroredVariable. |
| def assign_sub(self, value, use_locking=False, name=None, read_value=True): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.assign_sub(value, use_locking, name, read_value) |
| with ds_context.enter_or_assert_strategy(self._distribute_strategy): |
| if (ds_context.in_cross_replica_context() and |
| not values_util.in_replica_update_context()): |
| values_util.mark_as_unsaveable() |
| return values_util.on_read_assign_sub_cross_replica( |
| self, value, read_value=read_value) |
| else: |
| return super(SyncOnReadVariable, |
| self).assign_sub(value, use_locking, name, read_value) |
| |
| def assign_add(self, value, use_locking=False, name=None, read_value=True): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.assign_add(value, use_locking, name, read_value) |
| with ds_context.enter_or_assert_strategy(self._distribute_strategy): |
| if (ds_context.in_cross_replica_context() and |
| not values_util.in_replica_update_context()): |
| values_util.mark_as_unsaveable() |
| return values_util.on_read_assign_add_cross_replica( |
| self, value, read_value=read_value) |
| else: |
| return super(SyncOnReadVariable, |
| self).assign_add(value, use_locking, name, read_value) |
| |
| def assign(self, value, use_locking=False, name=None, read_value=True): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.assign(value, use_locking, name, read_value) |
| with ds_context.enter_or_assert_strategy(self._distribute_strategy): |
| if (ds_context.in_cross_replica_context() and |
| not values_util.in_replica_update_context()): |
| values_util.mark_as_unsaveable() |
| return values_util.on_read_assign_cross_replica( |
| self, value, read_value=read_value) |
| else: |
| return super(SyncOnReadVariable, |
| self).assign(value, use_locking, name, read_value) |
| |
| def _scatter_not_implemented(self, method): |
| raise NotImplementedError( |
| "Variables with `synchronization=ON_READ` doesn't support `%s`" % |
| method) |
| |
| def scatter_sub(self, *args, **kwargs): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_sub(*args, **kwargs) |
| self._scatter_not_implemented("scatter_sub") |
| |
| def scatter_add(self, *args, **kwargs): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_add(*args, **kwargs) |
| self._scatter_not_implemented("scatter_add") |
| |
| def scatter_mul(self, *args, **kwargs): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_mul(*args, **kwargs) |
| self._scatter_not_implemented("scatter_mul") |
| |
| def scatter_div(self, *args, **kwargs): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_div(*args, **kwargs) |
| self._scatter_not_implemented("scatter_div") |
| |
| def scatter_min(self, *args, **kwargs): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_min(*args, **kwargs) |
| self._scatter_not_implemented("scatter_min") |
| |
| def scatter_max(self, *args, **kwargs): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_max(*args, **kwargs) |
| self._scatter_not_implemented("scatter_max") |
| |
| def scatter_update(self, *args, **kwargs): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.scatter_update(*args, **kwargs) |
| self._scatter_not_implemented("scatter_update") |
| |
| def value(self): |
| if values_util.is_saving_non_distributed(): |
| return self._primary.value() |
| with ds_context.enter_or_assert_strategy(self._distribute_strategy): |
| if (ds_context.in_cross_replica_context() and |
| not values_util.in_replica_update_context()): |
| if self._aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA: |
| return self._get_replica(0).value() |
| return self._get_cross_replica() |
| else: |
| # _get_on_device_or_primary() returns a Variable. |
| return self._get_on_device_or_primary().value() |
| |
| def _get_cross_replica(self): |
| if self._aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA: |
| # Consider returning a tensor value here to make the return value of |
| # _get_cross_replica consistent. |
| return self._get_replica(0) |
| if self._aggregation == vs.VariableAggregation.SUM: |
| values_util.mark_as_unsaveable() |
| with ds_context.enter_or_assert_strategy(self._distribute_strategy): |
| return self._distribute_strategy.reduce( |
| reduce_util.ReduceOp.from_variable_aggregation(self._aggregation), |
| self, |
| axis=None) |
| |
| def _as_graph_element(self): |
| if values_util.is_saving_non_distributed(): |
| return self._primary._as_graph_element() # pylint: disable=protected-access |
| # pylint: disable=protected-access |
| with ds_context.enter_or_assert_strategy(self._distribute_strategy): |
| if ds_context.in_cross_replica_context(): |
| return ops.convert_to_tensor(self._get_cross_replica()) |
| return self._get()._as_graph_element() |
| |
| def _gather_saveables_for_checkpoint(self): |
| """Overrides Trackable method. |
| |
| This allows both name-based and object-based save and restore of |
| `SyncOnReadVariable`s. |
| |
| Returns: |
| A dictionary mapping attribute names to `SaveableObject` factories. |
| """ |
| |
| def _saveable_factory(name=self._common_name): |
| return _SyncOnReadSaveable(self, name) |
| |
| return {trackable.VARIABLE_VALUE_KEY: _saveable_factory} |
| |
| def _write_object_proto(self, proto, options): |
| """Update a SavedObject proto for the caller. |
| |
| If a DistributedVariable object supports this method, it will be called when |
| saving with a pre-built `SavedObject` proto representing the object, plus an |
| instance of `SaveOptions`. This method is then free to modify that proto |
| instance. |
| |
| `DistributedVariable` with `AUTO` or `ON_WRITE` synchronization optionally |
| write out information about their components to the |
| `experimental_distributed_variable_components` field of a |
| `SavedVariable` (depending on the `SaveOptions` variable policy). |
| |
| Args: |
| proto: A pre-built `SavedObject` proto for this object. It is assumed this |
| will be a `SavedVariable` instance. |
| options: A `SaveOptions` instance. |
| """ |
| pass |
| |
| |
| # Register a conversion functions which reads the value of the variable, |
| # allowing instances of the class to be used as tensors. |
| # DistributedVariable |
| def _tensor_conversion_distributed_var(var, dtype=None, name=None, |
| as_ref=False): |
| return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access |
| |
| |
| ops.register_tensor_conversion_function(DistributedVariable, |
| _tensor_conversion_distributed_var) |
| |
| |
| # MirroredVariables |
| def _tensor_conversion_mirrored(var, dtype=None, name=None, as_ref=False): |
| return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access |
| |
| |
| ops.register_tensor_conversion_function(MirroredVariable, |
| _tensor_conversion_mirrored) |
| |
| |
| # Mirrored Values |
| def _tensor_conversion_mirrored_val(value, dtype=None, name=None, as_ref=False): |
| return ops.convert_to_tensor( |
| value._get(), dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access |
| |
| |
| ops.register_tensor_conversion_function(Mirrored, |
| _tensor_conversion_mirrored_val) |
| |
| |
| # SyncOnReadVariables |
| def _tensor_conversion_sync_on_read(var, dtype=None, name=None, as_ref=False): |
| return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access |
| |
| |
| ops.register_tensor_conversion_function(SyncOnReadVariable, |
| _tensor_conversion_sync_on_read) |
| |
| |
| class VariablePolicy(object): |
| """Policy defining synchronization and aggregation of a distributed variable. |
| |
| Given `synchronization` and `aggregation` parameters set on a `tf.Variable` |
| during variable creation within `tf.distribute` scope, `tf.distribute` creates |
| an appropriate policy object and assigns it to the distributed variable. All |
| variable operations are delegated to the respective policy object. |
| """ |
| |
| def __init__(self, aggregation): |
| self._aggregation = aggregation |
| |
| def value(self): |
| raise NotImplementedError( |
| "This method should be overridden by sub-classes.") |
| |
| def _is_mirrored(self): |
| raise NotImplementedError( |
| "This method should be overridden by sub-classes.") |
| |
| def _as_graph_element(self, _): |
| raise NotImplementedError( |
| "This method should be overridden by sub-classes.") |
| |
| def _get_cross_replica(self, var): |
| raise NotImplementedError( |
| "This method should be overridden by sub-classes.") |
| |
| def _update_replica(self, var, update_fn, value, **kwargs): |
| raise NotImplementedError( |
| "This method should be overridden by sub-classes.") |
| |
| |
| class OnReadPolicy(VariablePolicy): |
| """Policy defined for `tf.VariableSynchronization.ON_READ` synchronization. |
| |
| This policy is created when `synchronization` is set to |
| `tf.VariableSynchronization.ON_READ` and `aggregation` is set to any of the |
| values allowed by the `tf.VariableAggregation` enum such as `NONE`, `SUM`, |
| `MEAN` or `ONLY_FIRST_REPLICA`when creating a `tf.Variable` in `tf.distribute` |
| scope. |
| """ |
| |
| def _is_mirrored(self): |
| return False |
| |
| def value(self, var): |
| with ds_context.enter_or_assert_strategy(var.distribute_strategy): |
| if (ds_context.in_cross_replica_context() and |
| not values_util.in_replica_update_context()): |
| if self._aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA: |
| return var._get_replica(0).value() # pylint: disable=protected-access |
| return var._get_cross_replica() # pylint: disable=protected-access |
| else: |
| return var._get_on_device_or_primary().value() # pylint: disable=protected-access |
| |
| def _as_graph_element(self, var): |
| with ds_context.enter_or_assert_strategy(var.distribute_strategy): |
| if ds_context.in_cross_replica_context(): |
| return ops.convert_to_tensor(var._get_cross_replica()) # pylint: disable=protected-access |
| return var._get()._as_graph_element() # pylint: disable=protected-access |
| |
| def _get_cross_replica(self, var): |
| if self._aggregation == vs.VariableAggregation.ONLY_FIRST_REPLICA: |
| return var._get_replica(0) # pylint: disable=protected-access |
| if self._aggregation == vs.VariableAggregation.SUM: |
| values_util.mark_as_unsaveable() |
| with ds_context.enter_or_assert_strategy(var.distribute_strategy): |
| return var.distribute_strategy.reduce( |
| reduce_util.ReduceOp.from_variable_aggregation(self._aggregation), |
| var, |
| axis=None) |
| |
| def _update_replica(self, var, update_fn, value, **kwargs): |
| return update_fn(var._get_on_device_or_primary(), value, **kwargs) # pylint: disable=protected-access |
| |
| def _scatter_not_implemented(self, method): |
| raise NotImplementedError( |
| "ON_READ variables doesn't support `%s` in cross replica context" % |
| method) |
| |
| def assign_sub(self, var, value, use_locking=False, name=None, |
| read_value=True): |
| """Subtracts a value from this variable.""" |
| with ds_context.enter_or_assert_strategy(var.distribute_strategy): |
| if (ds_context.in_cross_replica_context() and |
| not values_util.in_replica_update_context()): |
| values_util.mark_as_unsaveable() |
| return values_util.on_read_assign_sub_cross_replica( |
| var, value, read_value=read_value) |
| else: |
| return values_util.on_write_assign_sub( |
| var, value, use_locking=use_locking, name=name, |
| read_value=read_value) |
| |
| def assign_add(self, var, value, use_locking=False, name=None, |
| read_value=True): |
| """Adds a value to this variable.""" |
| with ds_context.enter_or_assert_strategy(var.distribute_strategy): |
| if (ds_context.in_cross_replica_context() and |
| not values_util.in_replica_update_context()): |
| values_util.mark_as_unsaveable() |
| return values_util.on_read_assign_add_cross_replica( |
| var, value, read_value=read_value) |
| else: |
| return values_util.on_write_assign_add( |
| var, value, use_locking=use_locking, name=name, |
| read_value=read_value) |
| |
| def assign(self, var, value, use_locking=False, name=None, read_value=True): |
| with ds_context.enter_or_assert_strategy(var.distribute_strategy): |
| if (ds_context.in_cross_replica_context() and |
| not values_util.in_replica_update_context()): |
| values_util.mark_as_unsaveable() |
| return values_util.on_read_assign_cross_replica(var, value, |
| read_value=read_value) |
| else: |
| return values_util.on_write_assign(var, value, |
| use_locking=use_locking, |
| name=name, |
| read_value=read_value) |
| |
| def scatter_sub(self, *args, **kwargs): |
| del args, kwargs |
| self._scatter_not_implemented("scatter_sub") |
| |
| def scatter_add(self, *args, **kwargs): |
| del args, kwargs |
| self._scatter_not_implemented("scatter_add") |
| |
| def scatter_mul(self, *args, **kwargs): |
| del args, kwargs |
| self._scatter_not_implemented("scatter_mul") |
| |
| def scatter_div(self, *args, **kwargs): |
| del args, kwargs |
| self._scatter_not_implemented("scatter_div") |
| |
| def scatter_min(self, *args, **kwargs): |
| del args, kwargs |
| self._scatter_not_implemented("scatter_min") |
| |
| def scatter_max(self, *args, **kwargs): |
| del args, kwargs |
| self._scatter_not_implemented("scatter_max") |
| |
| def scatter_update(self, *args, **kwargs): |
| del args, kwargs |
| self._scatter_not_implemented("scatter_update") |
| |
| def get_saveable(self, var, primary_var, name): |
| """Create a saveable object for the given variable.""" |
| return values_util.get_on_read_saveable(var, primary_var, name) |
| |
| def get_restore_ops(self, var, tensor): |
| """Restore the same value into all variables.""" |
| return values_util.get_on_read_restore_ops(var, tensor, self._aggregation) |
| |
| |
| class OnWritePolicy(VariablePolicy): |
| """Policy defined for `tf.VariableSynchronization.ON_WRITE` synchronization. |
| |
| This policy is created when the following `synchronization` and `aggregation` |
| parameters are specified when creating a `tf.Variable` in `tf.distribute` |
| scope and `synchronization` is equal to `tf.VariableSynchronization.ON_WRITE` |
| or `tf.VariableSynchronization.AUTO`. |
| """ |
| |
| def _is_mirrored(self): |
| return True |
| |
| def value(self, var): |
| return var._get_on_device_or_primary().value() # pylint: disable=protected-access |
| |
| def _as_graph_element(self, var): |
| return var._get_on_device_or_primary()._as_graph_element() # pylint: disable=protected-access |
| |
| def _get_cross_replica(self, var): |
| # Return identity, to avoid directly exposing the variable to the user and |
| # allowing it to be modified by mistake. |
| return array_ops.identity(var._get_on_device_or_primary()) # pylint: disable=protected-access |
| |
| def _update_replica(self, var, update_fn, value, **kwargs): |
| if var.aggregation == variables_lib.VariableAggregation.NONE: |
| return update_fn(var._get_on_device_or_primary(), value, **kwargs) # pylint: disable=protected-access |
| return _on_write_update_replica(var, update_fn, value, **kwargs) |
| |
| def assign(self, var, value, use_locking=False, name=None, read_value=True): |
| return values_util.on_write_assign(var, value, use_locking=use_locking, |
| name=name, read_value=read_value) |
| |
| def assign_add(self, var, value, use_locking=False, name=None, |
| read_value=True): |
| return values_util.on_write_assign_add(var, value, use_locking=use_locking, |
| name=name, read_value=read_value) |
| |
| def assign_sub(self, var, value, use_locking=False, name=None, |
| read_value=True): |
| return values_util.on_write_assign_sub(var, value, use_locking=use_locking, |
| name=name, read_value=read_value) |
| |
| def scatter_sub(self, var, sparse_delta, use_locking=False, name=None): |
| return values_util.scatter_sub(var, sparse_delta, use_locking=use_locking, |
| name=name) |
| |
| def scatter_add(self, var, sparse_delta, use_locking=False, name=None): |
| return values_util.scatter_add(var, sparse_delta, use_locking=use_locking, |
| name=name) |
| |
| def scatter_mul(self, var, sparse_delta, use_locking=False, name=None): |
| return values_util.scatter_mul(var, sparse_delta, use_locking=use_locking, |
| name=name) |
| |
| def scatter_div(self, var, sparse_delta, use_locking=False, name=None): |
| return values_util.scatter_div(var, sparse_delta, use_locking=use_locking, |
| name=name) |
| |
| def scatter_min(self, var, sparse_delta, use_locking=False, name=None): |
| if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and |
| self._aggregation != vs.VariableAggregation.NONE): |
| raise NotImplementedError(values_util.scatter_error_msg.format( |
| op_name="scatter_min", aggregation=self._aggregation)) |
| return values_util.scatter_min(var, sparse_delta, use_locking=use_locking, |
| name=name) |
| |
| def scatter_max(self, var, sparse_delta, use_locking=False, name=None): |
| if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and |
| self._aggregation != vs.VariableAggregation.NONE): |
| raise NotImplementedError(values_util.scatter_error_msg.format( |
| op_name="scatter_max", aggregation=self._aggregation)) |
| return values_util.scatter_max(var, sparse_delta, use_locking=use_locking, |
| name=name) |
| |
| def scatter_update(self, var, sparse_delta, use_locking=False, name=None): |
| if (self._aggregation != vs.VariableAggregation.ONLY_FIRST_REPLICA and |
| self._aggregation != vs.VariableAggregation.NONE): |
| raise NotImplementedError(values_util.scatter_error_msg.format( |
| op_name="scatter_update", aggregation=self._aggregation)) |
| return values_util.scatter_update(var, sparse_delta, |
| use_locking=use_locking, |
| name=name) |
| |
| def get_saveable(self, var, primary_var, name): |
| """Saveable ops for AUTO variables.""" |
| return values_util.get_on_write_saveable(var, primary_var, name) |
| |
| def get_restore_ops(self, var, tensor): |
| return values_util.get_on_write_restore_ops(var, tensor) |
| |
| def _write_object_proto(self, var, proto, options): |
| """Update a SavedObject proto for the caller. |
| |
| If a DistributedVariable object supports this method, it will be called when |
| saving with a pre-built `SavedObject` proto representing the object, plus an |
| instance of `SaveOptions`. This method is then free to modify that proto |
| instance. |
| |
| `DistributedVariable` with `AUTO` or `ON_WRITE` synchronization optionally |
| write out information about their components to the |
| `experimental_distributed_variable_components` field of a |
| `SavedVariable` (depending on the `SaveOptions` variable policy). |
| |
| Args: |
| var : A DistributedVariable object |
| proto: A pre-built `SavedObject` proto for this object. It is assumed this |
| will be a `SavedVariable` instance. |
| options: A `SaveOptions` instance. |
| """ |
| values_util.write_object_proto(var, proto, options) |
