Computing the weight-update-ratio in Tensorflow

I'm searching for a way to compute the weight-update-ratio for optimizer steps in Tensorflow. The weight-update-ratio is defined as the update-scale divided by the variable scale in each step and can be used for inspecting network training.

Ideally I want a non-intrusive way to compute it in a single session run, but couldn't accomplish quite what I was looking for. Since the update-scale and parameter scale are independent of the train step, one needs to add explicit dependencies to the graph in order to graph variable-scale before and after the update step. Unfortunately, it seems that in TF dependencies can only be defined for new nodes, which further complicates the issue.

So far, the best I've come up with is a context manager for definining the necessary operations. Its used as follows

opt = tf.train.AdamOptimizer(1e0)
grads = tf.gradients(loss, tf.trainable_variables())
grads = list(zip(grads, tf.trainable_variables()))

with compute_weight_update_ratio('wur') as wur:
    train = opt.apply_gradients(grads_and_vars=grads)

# ...
with tf.Session() as sess:
    sess.run(wur.ratio)

The full code of compute_weight_update_ratio can be found below. What bugs me is that in the current state the weight-update-ratio (at least norm_before) is computed with every training step, but for performance reason I'd rather prefer to do it selectively (e.g only when summaries are computed).

Any ideas on how to improve?

@contextlib.contextmanager
def compute_weight_update_ratio(name, var_scope=None):
    '''Injects training to compute weight-update-ratio.

    The weight-update-ratio is computed as the update scale divided
    by the variable scale before the update and should be somewhere in the 
    range 1e-2 or 1e-3.

    Params
    ------
    name : str
        Operation name

    Kwargs
    ------
    var_scope : str, optional
        Name selection of variables to compute weight-update-ration for. Defaults to all. Regex supported.
    '''

    class WeightUpdateRatio:
        def __init__(self):
            self.num_train = len(tf.get_collection(tf.GraphKeys.TRAIN_OP))
            self.variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=var_scope)
            self.norm_before = tf.norm(self.variables, name='norm_before')

        def compute_ratio(self,):
            train_ops = tf.get_collection(tf.GraphKeys.TRAIN_OP)
            assert len(train_ops) > self.num_train, 'Missing training op'

            with tf.control_dependencies(train_ops[self.num_train:]):
                self.norm_after = tf.norm(self.variables, name='norm_after')

            absdiff = tf.abs(tf.subtract(self.norm_after, self.norm_before), name='absdiff')
            self.ratio = tf.divide(absdiff, self.norm_before, name=name)

    with tf.name_scope(name) as scope:
        try:
            wur = WeightUpdateRatio()

            with tf.control_dependencies([wur.norm_before]):
                yield wur
        finally:
            wur.compute_ratio()