6436259ca4542e030697c05a72c942aa9aa3ced8,sonnet/src/rmsprop.py,RMSProp,apply,#RMSProp#Any#Any#,90

Before Change

        momentum = tf.cast(self.momentum, update.dtype.base_dtype)
        epsilon = tf.cast(self.epsilon, update.dtype.base_dtype)

        if self.centered:
          tf.raw_ops.ResourceApplyCenteredRMSProp(
              var=parameter.handle,
              mg=mg.handle,
              ms=ms.handle,
              mom=mom.handle,
              lr=learning_rate,
              rho=decay,
              momentum=momentum,
              epsilon=epsilon,
              grad=update)
        else:
          tf.raw_ops.ResourceApplyRMSProp(
              var=parameter.handle,
              ms=ms.handle,
              mom=mom.handle,
              lr=learning_rate,
              rho=decay,
              momentum=momentum,
              epsilon=epsilon,
              grad=update)


class ReferenceRMSProp(base.Module):
  Reference version of the RMSProp module.

  See: http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf

After Change

        epsilon = tf.cast(self.epsilon, update.dtype.base_dtype)

        ms.assign(tf.square(update) * (1. - decay) + ms * decay)
        if self.centered:
          mg.assign(update * (1. - decay) + mg * decay)
          denominator = ms - mg + epsilon
        else:
          denominator = ms + epsilon
        mom.assign(momentum * mom + (
            learning_rate * update * tf.math.rsqrt(denominator)))
        parameter.assign_sub(mom)

Instances