484ffb3eae65a09c717a6c823dcc888d5750df6c,sonnet/src/momentum.py,Momentum,apply,#Momentum#Any#Any#,67

Before Change

    self._initialize(parameters)
    for update, parameter, momentum in zip(updates, parameters,
                                           self.accumulated_momentum):
      if update is not None:
        optimizer_utils.check_same_dtype(update, parameter)
        lr = tf.cast(self.learning_rate, update.dtype)
        mu = tf.cast(self.momentum, update.dtype)
        if isinstance(update, tf.IndexedSlices):
          update, indices = optimizer_utils.deduplicate_indexed_slices(
              update.values, update.indices)
          sparse_momentum_update = (mu * momentum.sparse_read(indices)) + update
          momentum.scatter_update(
              tf.IndexedSlices(sparse_momentum_update, indices))
          if self.use_nesterov:
            parameter.scatter_sub(
                tf.IndexedSlices(
                    (lr * update) + (lr * mu * sparse_momentum_update),
                    indices))
          else:
            parameter.scatter_sub(
                tf.IndexedSlices(lr * sparse_momentum_update, indices))
        else:
          momentum.assign((mu * momentum) + update)
          if self.use_nesterov:
            parameter.assign_sub((lr * update) + (lr * mu * momentum))
          else:
            parameter.assign_sub(lr * momentum)


class FastMomentum(base.Optimizer):
  SGD with Momentum module.

  def __init__(self,

After Change

        momentum = momentum_var.sparse_read(indices)

        // Compute and apply a sparse update to our parameter and state.
        update, momentum = momentum_update(update, learning_rate, mu, momentum,
                                           self.use_nesterov)
        momentum_var.scatter_update(tf.IndexedSlices(momentum, indices))
        param.scatter_sub(tf.IndexedSlices(update, indices))

Instances