在tensorflow 2.x中使用_梯度替换为，以记录tensorboard上的梯度

Question

我试图用tensorflow和keras在tensorboard上绘制一个梯度直方图

我知道keras.callbacks.tensorboard的with_grads参数已经过时，似乎必须使用自定义回调来绘制渐变。我遇到了this post。从这个答案看来，必要的代码行是：

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

但是答案并没有真正指定将这些行放在哪里，以及是否确实需要回调：我是否应该将tf.keras.callbacks.TensorBoard子类化，并将这些行放在函数on_epoch_end中？差不多

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)

或者我可以在回调之外调用它们，例如在调用model.fit(...)之前

除此之外，当调用这些线时，我得到一个关于由线loss = model(model.trainable_weights)引起的张量维数的错误：

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

（784，因为我在MNIST图像上训练NN，并且图像为28x28=784像素

当使用Keras API时，有人能帮助掌握在tensorboard上绘制渐变信息的必要步骤吗