2

I am trying to plot a histogram of the gradient on tensorboard using tensorflow and keras.

I understood with_grads parameter of keras.callbacks.tensorboard is outdated and it seems one has to use a custom callback to plot the gradient. I came across this post. From that answer it seems the necessary lines of code are:

with tf.GradientTape() as tape:
    loss = model(model.trainable_weights) 
tape.gradient(loss, model.trainable_weights)

but the answer does not really specify where to put these lines and if a callback is actually needed: should I subclass the tf.keras.callbacks.TensorBoard and put these lines into the function on_epoch_end? Something like

class GradientCalcCallback(keras.callbacks.TensorBoard):
    def __init__(self, model_logs_directory, histogram_freq=1, write_graph=True, write_images=True):
        super(GradientCalcCallback, self).
            __init__(log_dir=model_logs_directory, histogram_freq=histogram_freq,
                     write_graph=True, write_images=write_images, update_freq='epoch')
    
     def on_epoch_end(self, epoch, logs=None):
        super().on_epoch_end(epoch, logs)
        with tf.GradientTape() as tape:
            loss = self.model(self.model.trainable_weights) 
        tape.gradient(loss, self.model.trainable_weights)

or can I call them outside the callback, for example before calling model.fit(...)?

Apart from that, when calling those lines I get an error about the dimensions of the tensors caused by the line loss = model(model.trainable_weights):

Tensor's shape (784, 784) is not compatible with supplied shape (784, 1)

(784 because I am training a NN on MNIST images and the images are 28x28=784 pixels

Can someone help to get to grips with the necessary steps to plot gradient information on tensorboard when using Keras APIs?

roschach
  • 8,390
  • 14
  • 74
  • 124
  • `Write_Grads` was not implemented in `TF 2.x`. This is one of the highly expected feature request [that is still open](https://github.com/tensorflow/tensorflow/issues/31173). As a workaround, we need to import `TF 1.x` modules and use `write_grads` as shown in the following code. For implementation please refer [this](https://stackoverflow.com/a/64062465/14290681). Thanks! –  Nov 29 '20 at 15:32
  • @TFer2 I know that there is a feature request for that, but now it has been around for a long time. I was wondering if using `tf.summary` and `tf.GradientTape` I can log this information for example overriding `train_step` of `keras.Model` or on a callback. – roschach Nov 29 '20 at 16:09

0 Answers0