如何使用 Tensorboard 检测消失和爆炸梯度？

Posted 2023-03-27

【中文标题】如何使用 Tensorboard 检测消失和爆炸梯度？

【英文标题】：How to detect vanishing and exploding gradients with Tensorboard?

【发布时间】：2020-06-03 02:34:25

【问题描述】：

我有两个“子问题”

1) 鉴于目前 write_grads=True 在 Tensorboard callback 中已被弃用（根据 "un-deprecate write_grads for fit #31173"），我如何使用 Tensorboard 检测消失或爆炸梯度？

2) 我想我可能可以根据 Tensorboard 的 Distributions and Histograms 选项卡中的权重分布和直方图判断我的模型是否存在梯度消失的问题。我的问题是我没有可比较的参考框架。目前，我的偏见似乎在“移动”，但我无法判断我的内核权重（Conv2D 层）是否“移动”/“改变”“足够”。有人可以通过给出经验法则来帮助我在 Tensorboard 中直观地评估这一点吗？ IE。如果只有底部 25% 的内核权重在移动，那是否足够好/不够好？或者也许有人可以从 tensorBoard 中发布两个消失梯度与非消失梯度的参考图像。

这是我的直方图和分布，是否可以判断我的模型是否存在梯度消失的问题？ （为简洁起见省略了一些层）提前致谢。

【参考方案1】：

我目前面临同样的问题，并使用 Tensorboard 类似地解决了这个问题。

即使 write_grads 已被弃用，您仍然可以通过子类化 tf.keras.Model 类并在 train_step 方法中使用 gradient.Tape 手动计算梯度来设法记录网络每一层的梯度。

类似的东西对我有用

from tensorflow.keras import Model

class TrainWithCustomLogsModel(Model):

    def __init__(self, **kwargs):
        super(TrainWithCustomLogsModel, self).__init__(**kwargs)
        self.step = tf.Variable(0, dtype=tf.int64,trainable=False)

    def train_step(self, data):

        # Get batch images and labels
        x, y = data
        
        # Compute the batch loss
        with tf.GradientTape() as tape:
            p = self(x , training = True)
            loss = self.compiled_loss(y, p, regularization_losses=self.losses)
        
        # Compute gradients for each weight of the network. Note trainable_vars and gradients are list of tensors
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        # Log gradients in Tensorboard
        self.step.assign_add(tf.constant(1, dtype=tf.int64))
        #tf.print(self.step)
        with train_summary_writer.as_default():
          for var, grad in zip(trainable_vars, gradients):
            name = var.name
            var, grad = tf.squeeze(var), tf.squeeze(grad)
            tf.summary.histogram(name, var, step = self.step)
            tf.summary.histogram('Gradients_'+name, grad, step = self.step)
    
        # Update model's weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        del tape
        # Update metrics (includes the metric that tracks the loss)
        self.compiled_metrics.update_state(y, p)
        # Return a dict mapping metric names to current value
        return m.name: m.result() for m in self.metrics

然后，您应该能够可视化训练的任何训练步骤的梯度分布，以及内核值的分布。

此外，可能值得尝试通过时间而不是单个值绘制范数分布。