Получите градиенты с помощью Keras Tensorflow 2.0

Question

Я хотел бы отслеживать градиенты на тензорной доске. Однако, поскольку сессия запуска заявление не вещь больше и write_grads аргумент tf.keras.callbacks.TensorBoard является depricated, я хотел бы знать, как следить за градиентов во время тренировки с Keras или tensorflow 2.0.

Мой текущий подход - создать для этой цели новый класс обратного вызова, но безуспешно. Может быть, кто-то еще знает, как делать такие продвинутые вещи.

Код, созданный для тестирования, показан ниже, но вызывает ошибки независимо от печати значения градиента на консоли или тензорной плате.

import tensorflow as tf
from tensorflow.python.keras import backend as K

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu', name='dense128'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax', name='dense10')
])

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])


class GradientCallback(tf.keras.callbacks.Callback):
    console = True

    def on_epoch_end(self, epoch, logs=None):
        weights = [w for w in self.model.trainable_weights if 'dense' in w.name and 'bias' in w.name]
        loss = self.model.total_loss
        optimizer = self.model.optimizer
        gradients = optimizer.get_gradients(loss, weights)
        for t in gradients:
            if self.console:
                print('Tensor: {}'.format(t.name))
                print('{}\n'.format(K.get_value(t)[:10]))
            else:
                tf.summary.histogram(t.name, data=t)


file_writer = tf.summary.create_file_writer("./metrics")
file_writer.set_as_default()

# write_grads has been removed
tensorboard_cb = tf.keras.callbacks.TensorBoard(histogram_freq=1, write_grads=True)
gradient_cb = GradientCallback()

model.fit(x_train, y_train, epochs=5, callbacks=[gradient_cb, tensorboard_cb])

• Вывод градиентов смещения на консоль (параметр консоли = True) приводит к: AttributeError: объект 'Tensor' не имеет атрибута 'numpy'
• Запись в тензорную плату (параметр консоли = False) создает:TypeError: Использованиеtf.Tensor как Python boolне допускается. Использоватьif t is not None: вместо того if t: чтобы проверить, определен ли тензор, и используйте операции TensorFlow, такие как tf.cond, для выполнения подграфов, обусловленных значением тензора.

Answer 1

Чтобы вычислить градиенты потерь относительно весов, используйте

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

tape.gradient(loss, model.trainable_weights)

Это (возможно, плохо) задокументировано на GradientTape.

Нам не нужно tape.watch переменная, потому что по умолчанию отслеживаются обучаемые параметры.

Как функцию это можно записать как

def gradient(model, x):
    x_tensor = tf.convert_to_tensor(x, dtype=tf.float32)
    with tf.GradientTape() as t:
        t.watch(x_tensor)
        loss = model(x_tensor)
    return t.gradient(loss, x_tensor).numpy()

Answer 2

Также посмотрите здесь: https://github.com/tensorflow/tensorflow/issues/31542#issuecomment-630495970

Ричардвт написал дочерний класс Tensorboard.

Я адаптировал его следующим образом:

class ExtendedTensorBoard(tf.keras.callbacks.TensorBoard):
    def _log_gradients(self, epoch):
        writer = self._get_writer(self._train_run_name)

        with writer.as_default(), tf.GradientTape() as g:
            # here we use test data to calculate the gradients
            features, y_true = list(val_dataset.batch(100).take(1))[0]

            y_pred = self.model(features)  # forward-propagation
            loss = self.model.compiled_loss(y_true=y_true, y_pred=y_pred)  # calculate loss
            gradients = g.gradient(loss, self.model.trainable_weights)  # back-propagation

            # In eager mode, grads does not have name, so we get names from model.trainable_weights
            for weights, grads in zip(self.model.trainable_weights, gradients):
                tf.summary.histogram(
                    weights.name.replace(':', '_') + '_grads', data=grads, step=epoch)

        writer.flush()

    def on_epoch_end(self, epoch, logs=None):
        # This function overwrites the on_epoch_end in tf.keras.callbacks.TensorBoard
        # but we do need to run the original on_epoch_end, so here we use the super function.
        super(ExtendedTensorBoard, self).on_epoch_end(epoch, logs=logs)

        if self.histogram_freq and epoch % self.histogram_freq == 0:
            self._log_gradients(epoch)