def apply_gradients(self, grads_and_vars, global_step=None, name=None):
"""Apply gradients to variables.
Args:
grads_and_vars: List of (gradient, variable) pairs as returned by
compute_gradients().
global_step: Optional Variable to increment by one after the
variables have been updated.
name: Optional name for the returned operation. Default to the
name passed to the Optimizer constructor.
Returns:
An `Operation` that applies the gradients. If `global_step` was not None,
that operation also increments `global_step`.
Raises:
ValueError: If the grads_and_vars is malformed.
"""
summed_grads_and_vars = []
for (grad, var) in grads_and_vars:
if grad is None:
summed_grads_and_vars.append((grad, var))
else:
with ops.colocate_with(grad):
summed_grads_and_vars.append(
(gen_poputil_ops.ipu_stateful_gradient_accumulate(
grad,
num_mini_batches=self._num_mini_batches,
verify_usage=self._verify_usage), var))
return self._opt.apply_gradients(summed_grads_and_vars, global_step,
name)
def apply_gradients(self, grads_and_vars, global_step=None, name=None):
summed_grads_and_vars = []
for (grad, var) in grads_and_vars:
if grad is None:
summed_grads_and_vars.append((grad, var))
else:
with ops.colocate_with(grad):
# gradient accumulation
if self._gradients_to_accumulate > 1 and not self._pipelining:
grad = gen_poputil_ops.ipu_stateful_gradient_accumulate(
grad / self._gradients_to_accumulate,
num_mini_batches=self._gradients_to_accumulate)
# replication
if self._replicas > 1:
grad = gen_poputil_ops.ipu_replication_normalise(
cross_replica_ops.cross_replica_sum(grad))
grad = math_ops.cast(grad, var.dtype)
summed_grads_and_vars.append((grad, var))
if self._pipelining:
# can do weight decay here as apply_gradients is only called on last accumulation step
summed_grads_and_vars = self.add_WD(summed_grads_and_vars)
ret = self._optimizer.apply_gradients(summed_grads_and_vars,
global_step, name)
if self._sharded:
sharding.propagate_sharding(ops.get_default_graph())
return ret
def body(i, y):
ga = gen_poputil_ops.ipu_stateful_gradient_accumulate(
array_ops.ones_like(y), num_mini_batches=5)
cr = gen_popops_ops.ipu_cross_replica_sum(ga)
y = y + cr
i = i + 1
return (i, y)
def body(i, x, y):
x = x + gen_poputil_ops.ipu_stateful_gradient_accumulate(
array_ops.ones_like(x),
num_mini_batches=5,
verify_usage=False)
y = y + array_ops.ones_like(x)
i = i + 1
return (i, x, y)
def body(i, y):
cr = gen_popops_ops.ipu_cross_replica_sum(
array_ops.ones_like(y))
norm = gen_poputil_ops.ipu_replication_normalise(cr)
ga = gen_poputil_ops.ipu_stateful_gradient_accumulate(
norm, num_mini_batches=5)
y = y + ga
i = i + 1
return (i, y)
def apply_gradients(self, grads_and_vars, global_step=None, name=None):
summed_grads_and_vars = []
for (grad, var) in grads_and_vars:
if grad is None:
summed_grads_and_vars.append((grad, var))
else:
with ops.colocate_with(grad):
# gradient accumulation
if self._gradient_accumulation_count > 1 and not self._pipelining:
grad = gen_poputil_ops.ipu_stateful_gradient_accumulate(
grad,
num_mini_batches=self._gradient_accumulation_count)
# replication
if self._replicas > 1:
grad = gen_poputil_ops.ipu_replication_normalise(
cross_replica_ops.cross_replica_sum(grad))
# distribution with IPUMultiWorkerStrategy needs additional normalisation by the number of workers
if isinstance(
distribute.get_strategy(),
ipu_multi_worker_strategy.IPUMultiWorkerStrategy):
grad /= distribute.get_strategy().num_replicas_in_sync
grad = math_ops.cast(grad, var.dtype)
summed_grads_and_vars.append((grad, var))
if self._pipelining:
# can do weight decay here as apply_gradients is only called on last accumulation step
summed_grads_and_vars = self.add_WD(summed_grads_and_vars)
if self._grad_scale != 1.0:
# don't rescale batch norm moving average statistics as they are not affected by loss scaling
summed_grads_and_vars = [
(grad, var) if 'batch_norm/moving_' in var.name else
(grad / self._grad_scale, var)
for grad, var in summed_grads_and_vars
]
ret = self._optimizer.apply_gradients(summed_grads_and_vars,
global_step, name)
if self._sharded:
sharding.propagate_sharding(ops.get_default_graph())
return ret
