class _DistributedOptimizer(_LegacyOptimizer):
"""An optimizer that wraps another tf.Optimizer, using an allreduce to
combine gradient values before applying gradients to model weights."""
def __init__(self,
optimizer,
name=None,
use_locking=False,
device_dense='',
device_sparse='',
compression=Compression.none,
sparse_as_dense=False,
op=Average,
gradient_predivide_factor=1.0,
backward_passes_per_step=1,
average_aggregated_gradients=False,
num_groups=0):
if name is None:
name = "Distributed{}".format(type(optimizer).__name__)
super(_DistributedOptimizer,
self).__init__(name=name, use_locking=use_locking)
self._optimizer = optimizer
self._allreduce_grads = _make_allreduce_grads_fn(
name, device_dense, device_sparse, compression,
sparse_as_dense, op, gradient_predivide_factor, num_groups)
self._agg_helper = None
if backward_passes_per_step > 1:
if _executing_eagerly():
raise ValueError(
"backward_passes_per_step > 1 is not yet supported "
"for _LegacyOptimizer with eager execution.")
self._agg_helper = LocalGradientAggregationHelper(
backward_passes_per_step=backward_passes_per_step,
allreduce_func=self._allreduce_grads,
sparse_as_dense=sparse_as_dense,
average_aggregated_gradients=average_aggregated_gradients,
rank=rank(),
optimizer_type=LocalGradientAggregationHelper.
_OPTIMIZER_TYPE_LEGACY,
)
def compute_gradients(self, *args, **kwargs):
"""Compute gradients of all trainable variables.
See Optimizer.compute_gradients() for more info.
In DistributedOptimizer, compute_gradients() is overriden to also
allreduce the gradients before returning them.
"""
gradients = self._optimizer.compute_gradients(*args, **kwargs)
grads, vars = zip(*gradients)
if self._agg_helper:
avg_grads = self._agg_helper.compute_gradients(grads)
else:
avg_grads = self._allreduce_grads(grads)
return list(zip(avg_grads, vars))
def apply_gradients(self, *args, **kwargs):
"""Calls this same method on the underlying optimizer."""
if self._agg_helper:
return self._agg_helper.apply_gradients(
lambda: self._optimizer.apply_gradients(*args, **kwargs),
self._optimizer,
*args,
**kwargs,
)
return self._optimizer.apply_gradients(*args, **kwargs)
def get_slot(self, *args, **kwargs):
"""Calls this same method on the underlying optimizer."""
return self._optimizer.get_slot(*args, **kwargs)
def get_slot_names(self, *args, **kwargs):
"""Calls this same method on the underlying optimizer."""
return self._optimizer.get_slot_names(*args, **kwargs)
def variables(self, *args, **kwargs):
"""Calls this same method on the underlying optimizer."""
return self._optimizer.variables(*args, **kwargs)
class _DistributedOptimizer(keras.optimizers.Optimizer):
_HAS_AGGREGATE_GRAD = True
def __init__(self, **kwargs):
self._name = name or "Distributed%s" % self.__class__.__base__.__name__
self._aggregated_gradients = False
self._allreduce_grads = hvd._make_allreduce_grads_fn(
self._name, device_dense, device_sparse, compression,
sparse_as_dense, op, gradient_predivide_factor)
self._agg_helper = None
if backward_passes_per_step > 1:
if hvd._executing_eagerly():
self._agg_helper = LocalGradientAggregationHelperEager(
backward_passes_per_step=backward_passes_per_step,
allreduce_func=self._allreduce_grads,
sparse_as_dense=sparse_as_dense,
average_aggregated_gradients=
average_aggregated_gradients,
)
else:
self._agg_helper = LocalGradientAggregationHelper(
backward_passes_per_step=backward_passes_per_step,
allreduce_func=self._allreduce_grads,
sparse_as_dense=sparse_as_dense,
average_aggregated_gradients=
average_aggregated_gradients,
rank=rank(),
optimizer_type=LocalGradientAggregationHelper.
_OPTIMIZER_TYPE_KERAS,
)
super(self.__class__, self).__init__(**kwargs)
def get_gradients(self, loss, params):
"""
Compute gradients of all trainable variables.
See Optimizer.get_gradients() for more info.
In DistributedOptimizer, get_gradients() is overriden to also
allreduce the gradients before returning them.
"""
gradients = super(self.__class__, self).get_gradients(loss, params)
return self._allreduce(gradients)
def _aggregate_gradients(self, grads_and_vars):
grads, vars = list(zip(*grads_and_vars))
aggregated_grads = self._allreduce(grads)
if _PRE_TF_2_4_0:
# Prior to TF 2.4.0, this function was expected to return only a list of
# grads, not a list of (grad, var) tuples.
return aggregated_grads
return list(zip(aggregated_grads, vars))
def _allreduce(self, grads):
self._aggregated_gradients = True
if self._agg_helper:
return self._agg_helper.compute_gradients(tuple(grads))
else:
return self._allreduce_grads(grads)
def apply_gradients(self, *args, **kwargs):
if self._agg_helper:
results = self._agg_helper.apply_gradients(
lambda: super(self.__class__, self).apply_gradients(
*args, **kwargs),
self,
*args,
**kwargs,
)
else:
results = super(self.__class__,
self).apply_gradients(*args, **kwargs)
if not self._aggregated_gradients:
raise Exception(
'`apply_gradients()` was called without a call to '
'`get_gradients()` or `_aggregate_gradients`. If you\'re '
'using TensorFlow 2.0, please specify '
'`experimental_run_tf_function=False` in `compile()`.')
return results
class _DistributedOptimizer(keras.optimizers.Optimizer):
_HAS_AGGREGATE_GRAD = True
def __init__(self, **kwargs):
self._name = name or "Distributed%s" % self.__class__.__base__.__name__
self._aggregated_gradients = False
self._allreduce_grads = hvd._make_allreduce_grads_fn(
self._name,
device_dense,
device_sparse,
compression,
sparse_as_dense,
op,
gradient_predivide_factor,
groups)
self._agg_helper = None
if backward_passes_per_step > 1:
if hvd._executing_eagerly():
self._agg_helper = LocalGradientAggregationHelperEager(
backward_passes_per_step=backward_passes_per_step,
allreduce_func=self._allreduce_grads,
sparse_as_dense=sparse_as_dense,
average_aggregated_gradients=average_aggregated_gradients,
)
else:
self._agg_helper = LocalGradientAggregationHelper(
backward_passes_per_step=backward_passes_per_step,
allreduce_func=self._allreduce_grads,
sparse_as_dense=sparse_as_dense,
average_aggregated_gradients=average_aggregated_gradients,
rank=rank(),
optimizer_type=LocalGradientAggregationHelper._OPTIMIZER_TYPE_KERAS,
)
super(self.__class__, self).__init__(**kwargs)
def _compute_gradients(self, loss, var_list, grad_loss=None, tape=None):
"""
Compute gradients of all trainable variables.
See Optimizer.get_gradients() for more info.
In DistributedOptimizer, get_gradients() is overriden to also
allreduce the gradients before returning them.
"""
if _PRE_TF_2_4_0:
return super(self.__class__, self)._compute_gradients(
loss, var_list, grad_loss, tape)
tape = backprop.GradientTape() if tape is None else tape
grads_and_vars = super(self.__class__, self)._compute_gradients(
# pylint: disable=protected-access
loss,
var_list,
grad_loss,
tape=tape)
grads, weights = list(zip(*grads_and_vars))
allreduced_grads = self._allreduce(grads, weights)
return list(zip(allreduced_grads, weights))
def get_gradients(self, loss, params):
"""
Compute gradients of all trainable variables.
See Optimizer.get_gradients() for more info.
In DistributedOptimizer, get_gradients() is overriden to also
allreduce the gradients before returning them.
"""
gradients = super(self.__class__, self).get_gradients(loss, params)
return self._allreduce(gradients, params)
def _aggregate_gradients(self, grads_and_vars):
if _PRE_TF_2_4_0:
grads, vars = list(zip(*grads_and_vars))
aggregated_grads = self._allreduce(grads, vars)
return aggregated_grads
else:
return super(self.__class__, self)._aggregate_gradients(
grads_and_vars)
def _allreduce(self, grads, vars):
self._aggregated_gradients = True
if self._agg_helper:
return self._agg_helper.compute_gradients(tuple(grads), tuple(vars))
else:
return self._allreduce_grads(grads, vars)
def apply_gradients(self, *args, **kwargs):
if self._agg_helper:
if isinstance(args[0], zip):
# If grad_and_vars are passed in as a zip object
# convert to a list. This is necessary for TF2.4+
# b/c args[0] is used in both conditional branches
# inside _agg_helper.apply_gradients().
args = list(args)
args[0] = list(args[0])
args = tuple(args)
results = self._agg_helper.apply_gradients(
lambda: super(self.__class__, self).apply_gradients(*args, **kwargs),
self,
*args,
**kwargs,
)
else:
results = super(self.__class__, self).apply_gradients(*args, **kwargs)
if _PRE_TF_2_4_0 and not self._aggregated_gradients:
raise Exception('`apply_gradients()` was called without a call to '
'`get_gradients()` or `_aggregate_gradients`. If you\'re '
'using TensorFlow 2.0, please specify '
'`experimental_run_tf_function=False` in `compile()`.')
return results