Как добавить пользовательские сводки градиента в tf.keras.callbacks.TensorBoard

Question

Я тренирую модель tf.keras и хотел бы использовать обратный вызов TensorBoard для записи градиентов на диск во время обучения.В частности, для каждой обучаемой переменной я вычисляю среднюю евклидову норму ее градиента на данном шаге и записываю ее в буфер скалярных сводок, например, так:

def add_mean_norm_summary(variable, scope):
    with tf.name_scope(scope):
        mean_l2_norm = tf.norm(variable, ord='euclidean') / tf.to_float(tf.size(variable))
        return tf.summary.scalar('{}_mean_l2_norm'.format(scope), mean_l2_norm)

Я написал пользовательский tf.keras.optimizer..Optimizer, который наследует от существующего оптимизатора (в моем случае Адама) и переопределяет его метод get_updates() следующим образом:

def get_updates(self, loss, params):

    grads = self.get_gradients(loss, params)

    grad_summaries = []
    for grad, param in zip(grads, params):
        summary_name_scope = param.name.replace(':', '_') + '_gradient'
        grad_summaries.append(add_mean_norm_summary(grad, summary_name_scope))
    tf.summary.merge(grad_summaries, name='gradient_summaries')

Я надеялся, что смогу затем написать собственный обратный вызов TensorBoard, который наследуется от tf.keras.callbacks.TensorBoard, где я мог бы каким-то образом взять эти резюме и добавить их в средство записи файлов методом on_batch_end().До сих пор я был неудачным.Я пытался вызывать опцию объединенных аннотаций по имени из моего производного класса TensorBoard, но он говорит мне, что в графе такого опера нет.

Кто-нибудь знает, как это сделать правильно?