我遇到了一个梯度爆炸问题，尝试了几天都没能解决。我在TensorFlow中实现了一个自定义的消息传递图神经网络，用于从图数据中预测一个连续值。每个图都与一个目标值相关联。图中的每个节点由一个节点属性向量表示，节点之间的边由一个边属性向量表示。

在消息传递层中，节点属性以某种方式更新（例如，通过聚合其他节点/边的属性），并返回这些更新后的节点属性。

现在，我已经弄清楚了代码中梯度问题发生的位置。我有以下代码片段。

to_concat = [neighbors_mean, e]z = K.concatenate(to_concat, axis=-1)output = self.Net(z)

这里，neighbors_mean是形成具有边属性e的边的两个节点属性vi、vj的逐元素平均值。Net是一个单层前馈网络。使用这种方法，在批次大小为32的情况下，训练损失在大约30个epoch后突然跳到NaN。如果批次大小为128，梯度仍然在大约200个epoch后爆炸。

我发现，在这种情况下，梯度爆炸是因为边属性e。如果我不将neighbors_mean与e连接起来，而是使用下面的代码，就不会发生梯度爆炸。

output = self.Net(neighbors_mean)

此外，我可以通过将e通过sigmoid函数发送来避免梯度爆炸。但这会降低性能（最终的MAE），因为e中的值被非线性地映射到0-1范围内。请注意，使用整流线性单元（ReLU）而不是sigmoid不起作用。

to_concat = [neighbors_mean, tf.math.sigmoid(e)]z = K.concatenate(to_concat, axis=-1)output = self.Net(z)

顺便提一下，e携带一个与两个对应节点之间的距离相关的单一值，这个距离始终在0.5到4的范围内。e中没有大值或NaN。

我有一个自定义的损失函数来训练这个模型，但我发现这不是损失函数的问题（其他损失函数也导致了相同的问题）。下面是我的自定义损失函数。请注意，尽管这是一个单输出回归网络，但我的神经网络的最后一层有两个神经元，分别与预测的均值和log(sigma)相关联。

def robust_loss(y_true, y_pred):  """  计算标签和预测之间的鲁棒损失。  """  mean, sigma = tf.split(y_pred, 2, axis=-1)  # 尝试限制'sigma'为  sigma = tf.clip_by_value(sigma,-4,1.0)，但梯度仍然爆炸  loss =  np.sqrt(2.0) * K.abs(mean - y_true) * K.exp(-sigma)  + sigma  return K.mean(loss)

我基本上尝试了网上建议的所有方法来避免梯度爆炸。

1. 应用了梯度裁剪 – 使用Adam(lr, clipnorm=1, clipvalue=5)和tf.clip_by_global_norm(gradients, 1.0)
2. 我的目标变量总是经过缩放
3. 权重使用glorot_uniform分布进行初始化
4. 对权重应用了正则化
5. 尝试了更大的批次大小（直到256，尽管在某个点上延迟了梯度爆炸）
6. 尝试降低学习率

我在这里遗漏了什么？我肯定知道这与连接e有关。但鉴于0.5<e<4，为什么在这种情况下梯度会爆炸？这个特征e对我很重要。我还能做些什么来避免模型中的数值溢出？

回答：

多亏了这个酷炫的调试工具tf.debugging.check_numerics，我解决了这个问题。

我最初确定连接e是问题所在，然后意识到传递给e的值远大于与e连接的neighbors_mean中的值。一旦它们被连接并通过神经网络（我的代码中的Net()）发送，我观察到一些输出按百计，随着训练的进行慢慢达到千计。

这是一个问题，因为我在消息传递层中有一个softmax操作。请注意，softmax计算一个指数（e^xi/Σe^xj）。任何超过e⁷⁰⁹的值在Python中都会导致数值溢出。这会产生inf值，最终一切变成nan是我代码中的问题。因此，这技术上不是一个梯度爆炸问题，这就是为什么它不能通过梯度裁剪解决。

我是如何追踪这个问题的？

我在几个我认为会产生nan值的层/张量下放置了tf.debugging.check_numerics()片段。类似这样：

tf.debugging.check_numerics(layerN, "LayerN is producing nans!")

一旦层输出在训练过程中变成inf或nan，这就会产生一个InvalidArgumentError。

Traceback (most recent call last):  File "trainer.py", line 506, in <module>    worker.train_model()  File "trainer.py", line 211, in train_model    l, tmae = train_step(*batch)  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/def_function.py", line 828, in __call__    result = self._call(*args, **kwds)  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/def_function.py", line 855, in _call    return self._stateless_fn(*args, **kwds)  # pylint: disable=not-callable  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/function.py", line 2943, in __call__    filtered_flat_args, captured_inputs=graph_function.captured_inputs)  # pylint: disable=protected-access  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/function.py", line 1919, in _call_flat    ctx, args, cancellation_manager=cancellation_manager))  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/function.py", line 560, in call    ctx=ctx)  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute    inputs, attrs, num_outputs)tensorflow.python.framework.errors_impl.InvalidArgumentError:  LayerN is producing nans! : Tensor had NaN values

现在我们知道问题出在哪里了。

如何解决这个问题

我对神经网络权重应用了核约束，这些权重的输出被传递到softmax函数。

layers.Dense(x, name="layer1", kernel_regularizer=regularizers.l2(1e-6), kernel_constraint=min_max_norm(min_value=1e-30, max_value=1.0))

这应该确保所有权重都小于1，并且层不会产生大的输出。这解决了问题，没有降低性能。

或者，可以使用softmax函数的数值稳定实现。