Calculating gradient norm wrt weights with keras

Question

I am attempting to calculate the gradient norm with respect to the weights of a neural network with keras (as a diagnostic tool). Eventually, I want to create a callback for

Answer 1

Extending josteinb's comment, I'm sharing the version that I have used.

Basically same with the previous answer, but this version integrates norm computation into the usual training routine.

import keras.backend as K

# Get a "l2 norm of gradients" tensor
def get_gradient_norm(model):
    with K.name_scope('gradient_norm'):
        grads = K.gradients(model.total_loss, model.trainable_weights)
        norm = K.sqrt(sum([K.sum(K.square(g)) for g in grads]))
    return norm

# Build a model
model = Model(...)

# Compile the model
model.compile(
    loss="categorical_crossentropy",
    optimizer="adam",
    metrics=["categorical_accuracy"],
)

# Append the "l2 norm of gradients" tensor as a metric
model.metrics_names.append("gradient_norm")
model.metrics_tensors.append(get_gradient_norm(model))

# You can compute the norm within the usual training routine
loss, acc, gradient_norm = model.train_on_batch(batch, label)