Gradient clipping in HMC?

[Josh Chang]

Is there a easy way to clip gradients within HMC?

[Josh Chang]

To clarify, this would be for TF 2

[Pavel Sountsov]

You can use custom_gradient to implement arbitrary gradient-post processing:

def clip_gradients(fn, clip_value):

def wrapper(*args):

@tf.custom_gradient

def grad_wrapper(*flat_args):

with tf.GradientTape() as tape:

tape.watch(flat_args)

ret = fn(*tf.nest.pack_sequence_as(args, flat_args))

def grad_fn(*dy):

flat_grads = tape.gradient(ret, flat_args, output_gradients=dy)

flat_grads = tf.nest.map_structure(lambda g: tf.where(tf.math.is_finite(g), g, tf.zeros_like(g)), flat_grads)

return tf.clip_by_global_norm(flat_grads, clip_value)[0]

return ret, grad_fn

return grad_wrapper(*tf.nest.flatten(args))

return wrapper

clipped_target_log_prob_fn = clip_gradients(target_log_prob_fn, clip_value=3.)

--
You received this message because you are subscribed to the Google Groups "TensorFlow Probability" group.
To unsubscribe from this group and stop receiving emails from it, send an email to tfprobabilit...@tensorflow.org.
To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/tfprobability/b61490a5-a61c-47f6-ac49-fe7928e6d648%40tensorflow.org.

[Pavel Sountsov]

(sorry, I sneaked in handling of NaN gradients as well in there, you can delete that line if you don't need that)

[Josh Chang]

Do you have a version of this offhand for kwargs?This is to wrap around a function like

def unormalized_log_prob(**x):

 x['x'] = data

 return joint_distribution_named.log_prob(x)

I tried this but it doesn't work (sorry I'm still fairly new to tensorflow workings)

def clip_gradients_dict(fn, clip_value):

   def wrapper(**args):

       @tf.custom_gradient

       def grad_wrapper(**flat_args):

           with tf.GradientTape() as tape:

               tape.watch(flat_args)

               ret = fn(**tf.nest.pack_sequence_as(args, flat_args))

           def grad_fn(*dy):

               flat_grads = tape.gradient(ret, flat_args, output_gradients=dy)

               flat_grads = tf.nest.map_structure(lambda g: tf.where(

                   tf.math.is_finite(g), g, tf.zeros_like(g)), flat_grads)

               return tf.clip_by_global_norm(flat_grads, clip_value)[0]

           return ret, grad_fn

       return grad_wrapper(**tf.nest.flatten(args))

   return wrapper

On Friday, September 13, 2019 at 2:09:36 PM UTC-4, Pavel Sountsov wrote:

(sorry, I sneaked in handling of NaN gradients as well in there, you can delete that line if you don't need that)

On Fri, Sep 13, 2019 at 11:07 AM Pavel Sountsov <si...@google.com> wrote:

You can use custom_gradient to implement arbitrary gradient-post processing:

def clip_gradients(fn, clip_value):

def wrapper(*args):

@tf.custom_gradient

def grad_wrapper(*flat_args):

with tf.GradientTape() as tape:

tape.watch(flat_args)

ret = fn(*tf.nest.pack_sequence_as(args, flat_args))

def grad_fn(*dy):

flat_grads = tape.gradient(ret, flat_args, output_gradients=dy)

flat_grads = tf.nest.map_structure(lambda g: tf.where(tf.math.is_finite(g), g, tf.zeros_like(g)), flat_grads)

return tf.clip_by_global_norm(flat_grads, clip_value)[0]

return ret, grad_fn

return grad_wrapper(*tf.nest.flatten(args))

return wrapper

clipped_target_log_prob_fn = clip_gradients(target_log_prob_fn, clip_value=3.)

On Fri, Sep 13, 2019 at 10:45 AM Josh Chang <josh.col...@gmail.com> wrote:

To clarify, this would be for TF 2

On Friday, September 13, 2019 at 1:45:23 PM UTC-4, Josh Chang wrote:

Is there a easy way to clip gradients within HMC?

--
You received this message because you are subscribed to the Google Groups "TensorFlow Probability" group.

To unsubscribe from this group and stop receiving emails from it, send an email to tfprob...@tensorflow.org.

[Pavel Sountsov]

You almost got it, the main error is that 'grad_wrapper' must only take *args. Here's a version of the clipping wrapper, but now supporting both *args and **kwargs:

def clip_gradients(fn, clip_value):

  def wrapper(*args, **kwargs):

    @tf.custom_gradient

    def grad_wrapper(*flat_args_kwargs):

      with tf.GradientTape() as tape:

        tape.watch(flat_args_kwargs)

        new_args, new_kwargs = tf.nest.pack_sequence_as((args, kwargs),

                                                        flat_args_kwargs)

        ret = fn(*new_args, **new_kwargs)

      def grad_fn(*dy):

        flat_grads = tape.gradient(ret, flat_args_kwargs, output_gradients=dy)

        flat_grads = tf.nest.map_structure(

            lambda g: tf.where(tf.math.is_finite(g), g, tf.zeros_like(g)),

            flat_grads)

        return tf.clip_by_global_norm(flat_grads, clip_value)[0]

      return ret, grad_fn

    return grad_wrapper(*tf.nest.flatten((args, kwargs)))

  return wrapper

To unsubscribe from this group and stop receiving emails from it, send an email to tfprobabilit...@tensorflow.org.
To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/tfprobability/62a2b4ce-fa47-4325-aed5-2b6ba49c473d%40tensorflow.org.