# 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. # ============================================================================== """Class MirroredStrategy implementing tf.distribute.Strategy.""" import copy from tensorflow.python import tf2 from tensorflow.python.distribute import collective_util from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib from tensorflow.python.distribute import cross_device_utils from tensorflow.python.distribute import device_util from tensorflow.python.distribute import distribute_lib from tensorflow.python.distribute import distribute_utils from tensorflow.python.distribute import input_lib from tensorflow.python.distribute import input_util from tensorflow.python.distribute import mirrored_run from tensorflow.python.distribute import multi_worker_util from tensorflow.python.distribute import numpy_dataset from tensorflow.python.distribute import reduce_util from tensorflow.python.distribute import values from tensorflow.python.distribute import values_util from tensorflow.python.distribute.cluster_resolver import tfconfig_cluster_resolver from tensorflow.python.distribute.v1 import input_lib as input_lib_v1 from tensorflow.python.eager import context from tensorflow.python.eager import record from tensorflow.python.framework import config from tensorflow.python.framework import constant_op from tensorflow.python.framework import device as tf_device from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import control_flow_util from tensorflow.python.ops import while_loop from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import nest from tensorflow.python.util.tf_export import tf_export # TODO(josh11b): Replace asserts in this file with if ...: raise ... def _is_device_list_single_worker(devices): """Checks whether the devices list is for single or multi-worker. Args: devices: a list of device strings or tf.config.LogicalDevice objects, for either local or for remote devices. Returns: a boolean indicating whether these device strings are for local or for remote. Raises: ValueError: if device strings are not consistent. """ specs = [] for d in devices: name = d.name if isinstance(d, context.LogicalDevice) else d specs.append(tf_device.DeviceSpec.from_string(name)) num_workers = len({(d.job, d.task, d.replica) for d in specs}) all_local = all(d.job in (None, "localhost") for d in specs) any_local = any(d.job in (None, "localhost") for d in specs) if any_local and not all_local: raise ValueError("Local device should have only 'localhost' in the job " "field in device string. " "E.g. 'job:localhost' in " "/job:localhost/replica:0/task:0/device:CPU:0" "Devices cannot have mixed list of device strings " "containing both localhost and other job types such as " "worker, ps etc. ") if num_workers == 1 and not all_local: if any(d.task is None for d in specs): raise ValueError("Remote device string must have task specified." "E.g. 'task:0' in " "/job:worker/replica:0/task:0/device:CPU:0") return num_workers == 1 def _cluster_spec_to_device_list(cluster_spec, num_gpus_per_worker): """Returns a device list given a cluster spec.""" cluster_spec = multi_worker_util.normalize_cluster_spec(cluster_spec) devices = [] for task_type in ("chief", "worker"): for task_id in range(len(cluster_spec.as_dict().get(task_type, []))): if num_gpus_per_worker == 0: devices.append("/job:%s/task:%d/device:CPU:0" % (task_type, task_id)) else: devices.extend([ "/job:%s/task:%d/device:GPU:%i" % (task_type, task_id, gpu_id) for gpu_id in range(num_gpus_per_worker) ]) return devices def _group_device_list(devices): """Groups the devices list by task_type and task_id. Args: devices: a list of device strings for remote devices. Returns: a dict of list of device strings mapping from task_type to a list of devices for the task_type in the ascending order of task_id. """ assert not _is_device_list_single_worker(devices) device_dict = {} for d in devices: d_spec = tf_device.DeviceSpec.from_string(d) # Create an entry for the task_type. if d_spec.job not in device_dict: device_dict[d_spec.job] = [] # Fill the device list for task_type until it covers the task_id. while len(device_dict[d_spec.job]) <= d_spec.task: device_dict[d_spec.job].append([]) device_dict[d_spec.job][d_spec.task].append(d) return device_dict def _is_gpu_device(device): return tf_device.DeviceSpec.from_string(device).device_type == "GPU" def _infer_num_gpus_per_worker(devices): """Infers the number of GPUs on each worker. Currently to make multi-worker cross device ops work, we need all workers to have the same number of GPUs. Args: devices: a list of device strings, can be either local devices or remote devices. Returns: number of GPUs per worker. Raises: ValueError if workers have different number of GPUs or GPU indices are not consecutive and starting from 0. """ if _is_device_list_single_worker(devices): return sum(1 for d in devices if _is_gpu_device(d)) else: device_dict = _group_device_list(devices) num_gpus = None for _, devices_in_task in device_dict.items(): for device_in_task in devices_in_task: if num_gpus is None: num_gpus = sum(1 for d in device_in_task if _is_gpu_device(d)) # Verify other workers have the same number of GPUs. elif num_gpus != sum(1 for d in device_in_task if _is_gpu_device(d)): raise ValueError("All workers should have the same number of GPUs.") for d in device_in_task: d_spec = tf_device.DeviceSpec.from_string(d) if (d_spec.device_type == "GPU" and d_spec.device_index >= num_gpus): raise ValueError("GPU `device_index` on a worker should be " "consecutive and start from 0.") return num_gpus def all_local_devices(num_gpus=None): devices = config.list_logical_devices("GPU") if num_gpus is not None: devices = devices[:num_gpus] return devices or config.list_logical_devices("CPU") def all_devices(): devices = [] tfconfig = tfconfig_cluster_resolver.TFConfigClusterResolver() if tfconfig.cluster_spec().as_dict(): devices = _cluster_spec_to_device_list(tfconfig.cluster_spec(), context.num_gpus()) return devices if devices else all_local_devices() @tf_export("distribute.MirroredStrategy", v1=[]) # pylint: disable=g-classes-have-attributes class MirroredStrategy(distribute_lib.Strategy): """Synchronous training across multiple replicas on one machine. This strategy is typically used for training on one machine with multiple GPUs. For TPUs, use `tf.distribute.TPUStrategy`. To use `MirroredStrategy` with multiple workers, please refer to `tf.distribute.experimental.MultiWorkerMirroredStrategy`. For example, a variable created under a `MirroredStrategy` is a `MirroredVariable`. If no devices are specified in the constructor argument of the strategy then it will use all the available GPUs. If no GPUs are found, it will use the available CPUs. Note that TensorFlow treats all CPUs on a machine as a single device, and uses threads internally for parallelism. >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"]) >>> with strategy.scope(): ... x = tf.Variable(1.) >>> x MirroredVariable:{ 0: , 1: } While using distribution strategies, all the variable creation should be done within the strategy's scope. This will replicate the variables across all the replicas and keep them in sync using an all-reduce algorithm. Variables created inside a `MirroredStrategy` which is wrapped with a `tf.function` are still `MirroredVariables`. >>> x = [] >>> @tf.function # Wrap the function with tf.function. ... def create_variable(): ... if not x: ... x.append(tf.Variable(1.)) ... return x[0] >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"]) >>> with strategy.scope(): ... _ = create_variable() ... print(x[0]) MirroredVariable:{ 0: , 1: } `experimental_distribute_dataset` can be used to distribute the dataset across the replicas when writing your own training loop. If you are using `.fit` and `.compile` methods available in `tf.keras`, then `tf.keras` will handle the distribution for you. For example: ```python my_strategy = tf.distribute.MirroredStrategy() with my_strategy.scope(): @tf.function def distribute_train_epoch(dataset): def replica_fn(input): # process input and return result return result total_result = 0 for x in dataset: per_replica_result = my_strategy.run(replica_fn, args=(x,)) total_result += my_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_result, axis=None) return total_result dist_dataset = my_strategy.experimental_distribute_dataset(dataset) for _ in range(EPOCHS): train_result = distribute_train_epoch(dist_dataset) ``` Args: devices: a list of device strings such as `['/gpu:0', '/gpu:1']`. If `None`, all available GPUs are used. If no GPUs are found, CPU is used. cross_device_ops: optional, a descendant of `CrossDeviceOps`. If this is not set, `NcclAllReduce()` will be used by default. One would customize this if NCCL isn't available or if a special implementation that exploits the particular hardware is available. """ # Only set this in tests. _collective_key_base = 0 def __init__(self, devices=None, cross_device_ops=None): extended = MirroredExtended( self, devices=devices, cross_device_ops=cross_device_ops) super(MirroredStrategy, self).__init__(extended) distribute_lib.distribution_strategy_gauge.get_cell("V2").set( "MirroredStrategy") @tf_export(v1=["distribute.MirroredStrategy"]) class MirroredStrategyV1(distribute_lib.StrategyV1): # pylint: disable=g-missing-docstring __doc__ = MirroredStrategy.__doc__ # Only set this in tests. _collective_key_base = 0 def __init__(self, devices=None, cross_device_ops=None): extended = MirroredExtended( self, devices=devices, cross_device_ops=cross_device_ops) super(MirroredStrategyV1, self).__init__(extended) distribute_lib.distribution_strategy_gauge.get_cell("V1").set( "MirroredStrategy") # TODO(josh11b): Switch to V2 when we no longer need to support tf.compat.v1. class MirroredExtended(distribute_lib.StrategyExtendedV1): """Implementation of MirroredStrategy.""" def __init__(self, container_strategy, devices=None, cross_device_ops=None): super(MirroredExtended, self).__init__(container_strategy) if context.executing_eagerly(): if devices and not _is_device_list_single_worker(devices): raise RuntimeError("In-graph multi-worker training with " "`MirroredStrategy` is not supported in eager mode.") else: if ( tfconfig_cluster_resolver.TFConfigClusterResolver() .cluster_spec() .as_dict() ): # if you are executing in eager mode, only the single machine code # path is supported. logging.info("Initializing local devices since in-graph multi-worker " "training with `MirroredStrategy` is not supported in " "eager mode. TF_CONFIG will be ignored when " "when initializing `MirroredStrategy`.") devices = devices or all_local_devices() else: devices = devices or all_devices() assert devices, ("Got an empty `devices` list and unable to recognize " "any local devices.") self._collective_key_base = container_strategy._collective_key_base self._communication_options = collective_util.Options( implementation=collective_util.CommunicationImplementation.NCCL) self._cross_device_ops = cross_device_ops self._initialize_strategy(devices) # TODO(b/128995245): Enable last partial batch support in graph mode. if ops.executing_eagerly_outside_functions(): self.experimental_enable_get_next_as_optional = True # Flag to turn on VariablePolicy. self._use_var_policy = False def _use_merge_call(self): # We currently only disable merge_call when XLA is used to compile the `fn` # passed to `strategy.run` and all devices are GPU. return not control_flow_util.GraphOrParentsInXlaContext( ops.get_default_graph()) or not all( [_is_gpu_device(d) for d in self._devices]) def _initialize_strategy(self, devices): # The _initialize_strategy method is intended to be used by distribute # coordinator as well. assert devices, "Must specify at least one device." devices = tuple(device_util.resolve(d) for d in devices) assert len(set(devices)) == len(devices), ( "No duplicates allowed in `devices` argument: %s" % (devices,)) self._initialize_single_worker(devices) self._collective_ops = self._make_collective_ops_with_fallbacks() # If cross_device_ops is not provided, set it to collective op by default. if not self._cross_device_ops: self._cross_device_ops = self._collective_ops def _make_collective_ops_with_fallbacks(self): self._collective_keys = cross_device_utils.CollectiveKeys( group_key_start=1 + self._collective_key_base) if not ops.executing_eagerly_outside_functions() and any( "gpu" not in d.lower() for d in self._devices): # In TF1/Session, fall back to ReductionToOneDevice() if there are # non-GPU devices or virtual GPUs are used. return cross_device_ops_lib.ReductionToOneDevice() # Use ReductionToOneDevice() if mixed devices are used. if any("cpu" in d.lower() for d in self._devices) and any( "gpu" in d.lower() for d in self._devices): return cross_device_ops_lib.ReductionToOneDevice() if all("cpu" in d.lower() for d in self._devices): # Use RING collective ops if all devices are CPU. self._communication_options = collective_util.Options( implementation=collective_util.CommunicationImplementation.RING) else: physical_gpus = context.context().list_physical_devices(device_type="GPU") logical_gpus = context.context().list_logical_devices(device_type="GPU") # Use RING collective ops if virtual devices are used. if len(physical_gpus) < len(logical_gpus): self._communication_options = collective_util.Options( implementation=collective_util.CommunicationImplementation.RING) # If all devices are physical GPU, use NCCL implementation. return cross_device_ops_lib.CollectiveAllReduce( devices=self._devices, group_size=len(self._devices), options=self._communication_options, collective_keys=self._collective_keys) def _initialize_single_worker(self, devices): """Initializes the object for single-worker training.""" self._devices = tuple(device_util.canonicalize(d) for d in devices) self._input_workers_devices = ( (device_util.canonicalize("/device:CPU:0", devices[0]), devices),) self._host_input_device = numpy_dataset.SingleDevice( self._input_workers_devices[0][0]) device_spec = tf_device.DeviceSpec.from_string( self._input_workers_devices[0][0]) # Ensures when we enter strategy.scope() we use the correct default device if device_spec.job is not None and device_spec.job != "localhost": self._default_device = "/job:%s/replica:%d/task:%d" % ( device_spec.job, device_spec.replica, device_spec.task) logging.info("Using MirroredStrategy with devices %r", devices) def _initialize_multi_worker(self, devices): """Initializes the object for multi-worker training.""" device_dict = _group_device_list(devices) workers = [] worker_devices = [] for job in ("chief", "worker"): for task in range(len(device_dict.get(job, []))): worker = "/job:%s/task:%d" % (job, task) workers.append(worker) worker_devices.append((worker, device_dict[job][task])) # Setting `_default_device` will add a device scope in the # distribution.scope. We set the default device to the first worker. When # users specify device under distribution.scope by # with tf.device("/cpu:0"): # ... # their ops will end up on the cpu device of its first worker, e.g. # "/job:worker/task:0/device:CPU:0". Note this is not used in replica mode. self._default_device = workers[0] self._host_input_device = numpy_dataset.SingleDevice(workers[0]) self._devices = tuple(devices) self._input_workers_devices = worker_devices self._is_multi_worker_training = True if len(workers) > 1: # Grandfather usage in the legacy tests if they're configured properly. if (not isinstance(self._cross_device_ops, cross_device_ops_lib.ReductionToOneDevice) or self._cross_device_ops._num_between_graph_workers > 1): # pylint: disable=protected-access raise ValueError( "In-graph multi-worker training with `MirroredStrategy` is not " "supported.") self._inferred_cross_device_ops = self._cross_device_ops else: # TODO(yuefengz): make `select_cross_device_ops` work with device strings # containing job names. self._inferred_cross_device_ops = cross_device_ops_lib.NcclAllReduce() logging.info("Using MirroredStrategy with remote devices %r", devices) def _input_workers_with_options(self, options=None): if not options: return input_lib.InputWorkers(self._input_workers_devices) if (options.experimental_replication_mode == distribute_lib.InputReplicationMode.PER_REPLICA): if options.experimental_place_dataset_on_device: self._input_workers_devices = ( tuple( (device_util.canonicalize(d, d), (d,)) for d in self._devices)) else: self._input_workers_devices = ( tuple((device_util.canonicalize("/device:CPU:0", d), (d,)) for d in self._devices)) return input_lib.InputWorkers(self._input_workers_devices) else: if not options.experimental_fetch_to_device: return input_lib.InputWorkers([ (host_device, (host_device,) * len(compute_devices)) for host_device, compute_devices in self._input_workers_devices ]) else: return input_lib.InputWorkers(self._input_workers_devices) @property def _input_workers(self): return self._input_workers_with_options() def _get_variable_creator_initial_value(self, replica_id, device, primary_var, **kwargs): """Return the initial value for variables on a replica.""" if replica_id == 0: return kwargs["initial_value"] else: assert primary_var is not None assert device is not None assert kwargs is not None def initial_value_fn(): if context.executing_eagerly() or ops.inside_function(): init_value = primary_var.value() return array_ops.identity(init_value) else: with ops.device(device): init_value = primary_var.initial_value return array_ops.identity(init_value) return initial_value_fn def _create_variable(self, next_creator, **kwargs): """Create a mirrored variable. See `DistributionStrategy.scope`.""" colocate_with = kwargs.pop("colocate_with", None) if colocate_with is None: devices = self._devices elif isinstance(colocate_with, numpy_dataset.SingleDevice): with ops.device(colocate_with.device): return next_creator(**kwargs) else: devices = colocate_with._devices # pylint: disable=protected-access def _real_mirrored_creator(**kwargs): # pylint: disable=g-missing-docstring value_list = [] for i, d in enumerate(devices): with ops.device(d): kwargs["initial_value"] = self._get_variable_creator_initial_value( replica_id=i, device=d, primary_var=value_list[0] if value_list else None, **kwargs) if i > 0: # Give replicas meaningful distinct names: var0name = value_list[0].name.split(":")[0] # We append a / to variable names created on replicas with id > 0 to # ensure that we ignore the name scope and instead use the given # name as the absolute name of the variable. kwargs["name"] = "%s/replica_%d/" % (var0name, i) with context.device_policy(context.DEVICE_PLACEMENT_SILENT): # Don't record operations (e.g. other variable reads) during # variable creation. with record.stop_recording(): v = next_creator(**kwargs) assert not isinstance(v, values.DistributedVariable) value_list.append(v) return value_list return distribute_utils.create_mirrored_variable( self._container_strategy(), _real_mirrored_creator, distribute_utils.VARIABLE_CLASS_MAPPING, distribute_utils.VARIABLE_POLICY_MAPPING, **kwargs) def _validate_colocate_with_variable(self, colocate_with_variable): distribute_utils.validate_colocate_distributed_variable( colocate_with_variable, self) def _make_dataset_iterator(self, dataset): return input_lib_v1.DatasetIterator( dataset, self._input_workers, self._container_strategy(), num_replicas_in_sync=self._num_replicas_in_sync) def _make_input_fn_iterator( self, input_fn, replication_mode=distribute_lib.InputReplicationMode.PER_WORKER): input_contexts = [] num_workers = self._input_workers.num_workers for i in range(num_workers): input_contexts.append(distribute_lib.InputContext( num_input_pipelines=num_workers, input_pipeline_id=i, num_replicas_in_sync=self._num_replicas_in_sync)) return input_lib_v1.InputFunctionIterator(input_fn, self._input_workers, input_contexts, self._container_strategy()) def _experimental_distribute_dataset(self, dataset, options): if (options and options.experimental_replication_mode == distribute_lib.InputReplicationMode.PER_REPLICA): raise NotImplementedError( "InputReplicationMode.PER_REPLICA " "is only supported in " "`distribute_datasets_from_function`." ) return input_util.get_distributed_dataset( dataset, self._input_workers_with_options(options), self._container_strategy(), num_replicas_in_sync=self._num_replicas_in_sync, options=options) def _experimental_make_numpy_dataset(self, numpy_input, session): return numpy_dataset.one_host_numpy_dataset( numpy_input, self._host_input_device, session) def _distribute_datasets_from_function(self, dataset_fn, options): input_workers = self._input_workers_with_options(options) input_contexts = [] num_workers = input_workers.num_workers for i in range(num_workers): input_contexts.append(distribute_lib.InputContext( num_input_pipelines=num_workers, input_pipeline_id=i, num_replicas_in_sync=self._num_replicas_in_sync)) return input_util.get_distributed_datasets_from_function( dataset_fn, input_workers, input_contexts, self._container_strategy(), options) def _experimental_distribute_values_from_function(self, value_fn): per_replica_values = [] for replica_id in range(self._num_replicas_in_sync): per_replica_values.append(value_fn( distribute_lib.ValueContext(replica_id, self._num_replicas_in_sync))) return distribute_utils.regroup(per_replica_values, always_wrap=True) # TODO(priyag): Deal with OutOfRange errors once b/111349762 is fixed. def _experimental_run_steps_on_iterator(self, fn, iterator, iterations, initial_loop_values=None): if initial_loop_values is None: initial_loop_values = {} initial_loop_values = nest.flatten(initial_loop_values) ctx = input_lib.MultiStepContext() def body(i, *args): """A wrapper around `fn` to create the while loop body.""" del args fn_result = fn(ctx, iterator.get_next()) for (name, output) in ctx.last_step_outputs.items(): # Convert all outputs to tensors, potentially from `DistributedValues`. ctx.last_step_outputs[name] = self._local_results(output) flat_last_step_outputs = nest.flatten(ctx.last_step_outputs) with ops.control_dependencies([fn_result]): return [i + 1] + flat_last_step_outputs # We capture the control_flow_context at this point, before we run `fn` # inside a while_loop. This is useful in cases where we might need to exit # these contexts and get back to the outer context to do some things, for # e.g. create an op which should be evaluated only once at the end of the # loop on the host. One such usage is in creating metrics' value op. self._outer_control_flow_context = ( ops.get_default_graph()._get_control_flow_context()) # pylint: disable=protected-access cond = lambda i, *args: i < iterations i = constant_op.constant(0) loop_result = while_loop.while_loop( cond, body, [i] + initial_loop_values, name="", parallel_iterations=1, back_prop=False, swap_memory=False, return_same_structure=True) del self._outer_control_flow_context ctx.run_op = control_flow_ops.group(loop_result) # Convert the last_step_outputs from a list to the original dict structure # of last_step_outputs. last_step_tensor_outputs = loop_result[1:] last_step_tensor_outputs_dict = nest.pack_sequence_as( ctx.last_step_outputs, last_step_tensor_outputs) for name, reduce_op in ctx._last_step_outputs_reduce_ops.items(): # pylint: disable=protected-access output = last_step_tensor_outputs_dict[name] # For outputs that have already been reduced, wrap them in a Mirrored # container, else in a PerReplica container. if reduce_op is None: last_step_tensor_outputs_dict[name] = distribute_utils.regroup(output) else: assert len(output) == 1 last_step_tensor_outputs_dict[name] = output[0] ctx._set_last_step_outputs(last_step_tensor_outputs_dict) # pylint: disable=protected-access return ctx def _broadcast_to(self, tensor, destinations): # This is both a fast path for Python constants, and a way to delay # converting Python values to a tensor until we know what type it # should be converted to. Otherwise we have trouble with: # global_step.assign_add(1) # since the `1` gets broadcast as an int32 but global_step is int64. if isinstance(tensor, (float, int)): return tensor # TODO(josh11b): In eager mode, use one thread per device, or async mode. if not destinations: # TODO(josh11b): Use current logical device instead of 0 here. destinations = self._devices return self._get_cross_device_ops(tensor).broadcast(tensor, destinations) def _call_for_each_replica(self, fn, args, kwargs): return mirrored_run.call_for_each_replica( self._container_strategy(), fn, args, kwargs) def _configure(self, session_config=None, cluster_spec=None, task_type=None, task_id=None): del task_type, task_id if session_config: session_config.CopyFrom(self._update_config_proto(session_config)) if cluster_spec: # TODO(yuefengz): remove the following code once cluster_resolver is # added. num_gpus_per_worker = _infer_num_gpus_per_worker(self._devices) multi_worker_devices = _cluster_spec_to_device_list( cluster_spec, num_gpus_per_worker) self._initialize_multi_worker(multi_worker_devices) def _update_config_proto(self, config_proto): updated_config = copy.deepcopy(config_proto) updated_config.isolate_session_state = True return updated_config def _get_cross_device_ops(self, value): # Always use CollectiveAllReduce when XLA is enabled, since other cross # device ops don't have as good support on XLA. if not self._use_merge_call(): if not isinstance(self._cross_device_ops, cross_device_ops_lib.CollectiveAllReduce): logging.warning( "Under XLA context, MirroredStrategy uses CollectiveAllReduce op. " "Although %r is provided to initialize MirroredStrategy, it is " "ignored in XLA. Please use CollectiveAllReduce(or default option) " "in the future, since other cross device ops are not well " "supported on XLA.", self._cross_device_ops ) return self._collective_ops if isinstance(value, values.DistributedValues): value_int32 = True in { dtypes.as_dtype(v.dtype) == dtypes.int32 for v in value.values } else: value_int32 = dtypes.as_dtype(value.dtype) == dtypes.int32 if value_int32: return cross_device_ops_lib.ReductionToOneDevice() else: return self._cross_device_ops def _gather_to_implementation(self, value, destinations, axis, options): if not isinstance(value, values.DistributedValues): # ReductionToOneDevice._gather accepts DistributedValues only. return value return self._get_cross_device_ops(value)._gather( # pylint: disable=protected-access value, destinations=destinations, axis=axis, options=self._communication_options.merge(options)) def _reduce_to(self, reduce_op, value, destinations, options): if (distribute_utils.is_mirrored(value) and reduce_op == reduce_util.ReduceOp.MEAN): return value assert not distribute_utils.is_mirrored(value) def get_values(value): if not isinstance(value, values.DistributedValues): # This function handles reducing values that are not PerReplica or # Mirrored values. For example, the same value could be present on all # replicas in which case `value` would be a single value or value could # be 0. return cross_device_ops_lib.reduce_non_distributed_value( reduce_op, value, destinations, self._num_replicas_in_sync) if self._use_merge_call() and ( not cross_device_ops_lib._devices_match(value, destinations) or # pylint: disable=protected-access any("cpu" in d.lower() for d in cross_device_ops_lib.get_devices_from(destinations))): return cross_device_ops_lib.ReductionToOneDevice().reduce( reduce_op, value, destinations) return self._get_cross_device_ops(value).reduce( reduce_op, value, destinations=destinations, options=self._communication_options.merge(options)) return nest.map_structure(get_values, value) def _batch_reduce_to(self, reduce_op, value_destination_pairs, options): cross_device_ops = None for value, _ in value_destination_pairs: if cross_device_ops is None: cross_device_ops = self._get_cross_device_ops(value) elif cross_device_ops is not self._get_cross_device_ops(value): raise ValueError("Inputs to batch_reduce_to must be either all on " "the host or all on the compute devices.") return cross_device_ops.batch_reduce( reduce_op, value_destination_pairs, options=self._communication_options.merge(options)) def _update(self, var, fn, args, kwargs, group): # TODO(josh11b): In eager mode, use one thread per device. assert isinstance(var, values.DistributedVariable) updates = [] for i, v in enumerate(var.values): name = "update_%d" % i with ops.device(v.device), \ distribute_lib.UpdateContext(i), \ ops.name_scope(name): # If args and kwargs are not mirrored, the value is returned as is. updates.append( fn(v, *distribute_utils.select_replica(i, args), **distribute_utils.select_replica(i, kwargs))) return distribute_utils.update_regroup(self, updates, group) def _replica_ctx_all_reduce(self, reduce_op, value, options=None): """Implements `StrategyExtendedV2._replica_ctx_all_reduce`.""" # This implementation avoids using `merge_call` and just launches collective # ops in one replica. if options is None: options = collective_util.Options() if context.executing_eagerly() or ( not tf2.enabled()) or self._use_merge_call(): # In eager mode, falls back to the default implementation that uses # `merge_call`. Replica functions are running sequentially in eager mode, # and due to the blocking nature of collective ops, execution will hang if # collective ops are to be launched sequentially. return super()._replica_ctx_all_reduce(reduce_op, value, options) replica_context = distribute_lib.get_replica_context() assert replica_context, ( "`StrategyExtended._replica_ctx_all_reduce` must be called in a " "replica context") return self._get_cross_device_ops(value)._all_reduce( # pylint: disable=protected-access reduce_op, value, replica_context._replica_id, # pylint: disable=protected-access options) def _replica_ctx_update(self, var, fn, args, kwargs, group): if self._use_merge_call(): return super()._replica_ctx_update(var, fn, args, kwargs, group) replica_context = distribute_lib.get_replica_context() assert replica_context replica_id = values_util.get_current_replica_id_as_int() name = "update_%d" % replica_id if isinstance(var, values.DistributedVariable): var = var._get_replica(replica_id) # pylint: disable=protected-access with ops.device(var.device), ops.name_scope(name): result = fn(var, *args, **kwargs) return result def _update_non_slot(self, colocate_with, fn, args, kwargs, group): assert isinstance(colocate_with, tuple) # TODO(josh11b): In eager mode, use one thread per device. updates = [] for i, d in enumerate(colocate_with): name = "update_%d" % i with ops.device(d), distribute_lib.UpdateContext(i), ops.name_scope(name): updates.append( fn(*distribute_utils.select_replica(i, args), **distribute_utils.select_replica(i, kwargs))) return distribute_utils.update_regroup(self, updates, group) def read_var(self, replica_local_var): """Read the aggregate value of a replica-local variable.""" # pylint: disable=protected-access if distribute_utils.is_sync_on_read(replica_local_var): return replica_local_var._get_cross_replica() assert distribute_utils.is_mirrored(replica_local_var) return array_ops.identity(replica_local_var._get()) # pylint: enable=protected-access def value_container(self, val): return distribute_utils.value_container(val) @property def _num_replicas_in_sync(self): return len(self._devices) @property def worker_devices(self): return self._devices @property def worker_devices_by_replica(self): return [[d] for d in self._devices] @property def parameter_devices(self): return self.worker_devices @property def experimental_between_graph(self): return False @property def experimental_should_init(self): return True @property def should_checkpoint(self): return True @property def should_save_summary(self): return True def non_slot_devices(self, var_list): del var_list # TODO(josh11b): Should this be the last logical device instead? return self._devices # TODO(priyag): Delete this once all strategies use global batch size. @property def _global_batch_size(self): """`make_dataset_iterator` and `make_numpy_iterator` use global batch size. `make_input_fn_iterator` assumes per-replica batching. Returns: Boolean. """ return True def _in_multi_worker_mode(self): """Whether this strategy indicates working in multi-worker settings.""" return False def _get_local_replica_id(self, replica_id_in_sync_group): return replica_id_in_sync_group def _get_replica_id_in_sync_group(self, replica_id): return replica_id