TensorFlow/Spektral图神经网络梯度下降无梯度问题排查

我在交互式学习场景下用图神经网络（GNN）做动作识别，通过动作值计算损失后执行梯度下降，但梯度下降环节出问题了——梯度函数返回全为None，导致权重更新失败。

之前在回归任务里用GNN执行梯度下降一切正常，现在用Spektral库结合TensorFlow imperative风格写DQN，就出现了这个问题。下面是可复现的代码和错误信息：

可复现代码 #

class GIN0(Model):
    def __init__(self, channels, n_layers):
        super().__init__()
        self.conv1 = GINConv(channels, epsilon=0, mlp_hidden=[channels, channels])
        self.convs = []
        for _ in range(1, n_layers):
            self.convs.append(
                GINConv(channels, epsilon=0, mlp_hidden=[channels, channels])
            )
        self.pool = GlobalAvgPool()
        self.dense1 = Dense(channels, activation="relu")
        self.dropout = Dropout(0.5)
        self.dense2 = Dense(channels, activation="relu")

    def call(self, inputs):
        x, a, i = inputs
        x = self.conv1([x, a])
        for conv in self.convs:
            x = conv([x, a])
        x = self.pool([x, i])
        x = self.dense1(x)
        x = self.dropout(x)
        return self.dense2(x)
class IGDQN(object):
    def __init__(self,
                 number_of_outputs,
                 layers,
                 alpha,
                 gamma,
                 epsilon
        ):
        self.number_of_outputs = number_of_outputs
        self.layers = layers
        self.alpha = alpha
        self.gamma = gamma
        self.epsilon = epsilon
        self.opt = Adam(lr=alpha)
        self.model = GIN0(number_of_outputs, layers)

    def choose_action(self, state, debug=False):
        if np.random.rand() < self.epsilon:
            return random.randrange(self.number_of_outputs)
        q = self.model.predict(state)
        if debug:
            print('q=',q)
            print('action_code=',np.argmin(q[0]))
        return np.argmin(q[0])

    @tf.function
    def update(self, loss):
        with tf.GradientTape(persistent=True) as tape:
            #the gin0 network weights are updated
            gradients = tape.gradient(loss, self.model.trainable_variables)
            print(gradients)
            self.opt.apply_gradients(zip(gradients, self.model.trainable_variables))

def get_inputs():
    indices = [
     [0, 1],
     [0, 2],
     [0, 4],
     [1, 0],
     [1, 2],
     [1, 3],
     [1, 5],
     [2, 0],
     [2, 1],
     [2, 3],
     [2, 4],
     [3, 1],
     [3, 2],
     [3, 7],
     [4, 0],
     [4, 2],
     [4, 5],
     [4, 6],
     [5, 1],
     [5, 4],
     [5, 6],
     [6, 4],
     [6, 5],
     [6, 7],
     [7, 3],
     [7, 6]]
    values = [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]
    dense_shape = [8,8]
    adjacency_matrix = tf.sparse.SparseTensor(
        indices, values, dense_shape
    )
    matrix = [
        [0., 0., 0., 1., 0., 6., 1.,],
        [0., 0., 0., 1., 0., 7., 0.,],
        [0., 0., 0., 1., 0., 1., 2.,],
        [0., 0., 0., 1., 0., 1., 3.,],
        [0., 0., 0., 1., 0., 6., 0.,],
        [0., 0., 0., 1., 0., 7., 1.,],
        [0., 0., 0., 1., 0., 0., 3.,],
        [0., 0., 0., 1., 0., 0., 2.,],
    ]
    properties_matrix = np.array(matrix)
    am = tf.sparse.to_dense(adjacency_matrix)
    g = Graph( x=properties_matrix, a=am.numpy(), e=None,y=[456] )
    ds = [g]
    design_name = PLconfig_grid.designName
    dsr = CircuitDataset2(design_name, ds, False, path="/home/xx/CircuitAttributePrediction/dataset")
    loader = DisjointLoader(dsr, batch_size=1)
    inputs, target = loader.__next__()
    return inputs

def check_IGDQN(designName, inputDir):
    number_of_outputs = 128
    layers = 3
    alpha = 5e-4
    gamma = 0.2
    epsilon = 0.3
    dqn = IGDQN(
            number_of_outputs,
            layers,
            alpha,
            gamma,
            epsilon
    )

    inputs = get_inputs()
    next_state = state = inputs
    action = dqn.choose_action(state)
    #loss calculation steps simplified for debug purposes
    loss = tf.constant(100, dtype=tf.float32)
    dqn.update(loss)

错误信息 #

[None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None]
Traceback (most recent call last):
  File "test_PLKerasNetworks_GIN0.py", line 142, in <module>
    main()
  File "test_PLKerasNetworks_GIN0.py", line 136, in main
    check_IGDQN(designName, inputDir)    
  File "test_PLKerasNetworks_GIN0.py", line 130, in check_IGDQN
    dqn.update(loss)
  File "/home/xx/.local/share/virtualenvs/xx-TxBsk36Y/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 828, in __call__
    result = self._call(*args, **kwds)
  File "/home/xx/.local/share/virtualenvs/xx-TxBsk36Y/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 871, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "/home/xx/.local/share/virtualenvs/xx-TxBsk36Y/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 726, in _initialize
    *args, **kwds))
  File "/home/xx/.local/share/virtualenvs/xx-TxBsk36Y/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2969, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "/home/xx/.local/share/virtualenvs/xx-TxBsk36Y/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3361, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/home/xx/.local/share/virtualenvs/xx-TxBsk36Y/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3206, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/home/xx/.local/share/virtualenvs/xx-TxBsk36Y/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 990, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/home/xx/.local/share/virtualenvs/xx-TxBsk36Y/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 634, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/home/xx/.local/share/virtualenvs/xx-TxBsk36Y/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3887, in bound_method_wrapper
    return wrapped_fn(*args, **kwargs)
  File "/home/xx/.local/share/virtualenvs/xx-TxBsk36Y/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 977, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:

    test_PLKerasNetworks_GIN0.py:56 update  *
        self.opt.apply_gradients(zip(gradients, self.model.trainable_variables))
    /home/xx/.local/share/virtualenvs/xx-TxBsk36Y/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:598 apply_gradients  **
        grads_and_vars = optimizer_utils.filter_empty_gradients(grads_and_vars)
    /home/xx/.local/share/virtualenvs/xx-TxBsk36Y/lib/python3.7/site-packages/tensorflow/python/keras/optimizer_v2/utils.py:79 filter_empty_gradients
        ([v.name for _, v in grads_and_vars],))

    ValueError: No gradients provided for any variable: ['dense/kernel:0', 'dense/bias:0', 'dense_1/kernel:0', 'dense_1/bias:0', 'dense_2/kernel:0', 'dense_3/kernel:0', 'dense_3/bias:0', 'dense_4/kernel:0', 'dense_4/bias:0', 'dense_5/kernel:0', 'dense_6/kernel:0', 'dense_6/bias:0', 'dense_7/kernel:0', 'dense_7/bias:0', 'dense_8/kernel:0', 'gi_n0/dense/kernel:0', 'gi_n0/dense/bias:0', 'gi_n0/dense_1/kernel:0', 'gi_n0/dense_1/bias:0'].

解决方案 #

1. 核心问题：损失与模型输出无依赖 #

你传入update方法的loss是固定的tf.constant(100)，和模型的输出、权重完全没有计算关联，TensorFlow无法追踪到梯度，因此返回全None。在DQN流程中，损失必须基于模型对状态的预测值和目标Q值计算，建立和模型前向传播的依赖关系。

2. 代码修正要点 #

（1）重构update方法，在梯度带内执行前向传播

修改IGDQN的update方法，传入状态和目标Q值，在tf.GradientTape范围内执行模型前向计算，再计算损失：

@tf.function
def update(self, state, target_q):
    with tf.GradientTape() as tape:
        # 在梯度带内执行前向传播，让TensorFlow追踪依赖
        pred_q = self.model(state, training=True)
        # 用DQN常用的MSE损失，可根据任务调整
        loss = tf.reduce_mean(tf.square(pred_q - target_q))
    
    # 计算梯度并更新权重
    gradients = tape.gradient(loss, self.model.trainable_variables)
    self.opt.apply_gradients(zip(gradients, self.model.trainable_variables))
    return loss

（2）修正GIN模型的层注册问题

GIN0类中self.convs用了普通列表，TensorFlow无法自动追踪这些层的变量，需要改用tf.keras.layers.LayerList管理：

from tensorflow.keras.layers import LayerList

class GIN0(Model):
    def __init__(self, channels, n_layers):
        super().__init__()
        self.conv1 = GINConv(channels, epsilon=0, mlp_hidden=[channels, channels])
        # 用LayerList替代普通列表
        self.convs = LayerList()
        for _ in range(1, n_layers):
            self.convs.add(
                GINConv(channels, epsilon=0, mlp_hidden=[channels, channels])
            )
        self.pool = GlobalAvgPool()
        self.dense1 = Dense(channels, activation="relu")
        self.dropout = Dropout(0.5)
        self.dense2 = Dense(channels, activation="relu")
    
    def call(self, inputs):
        x, a, i = inputs
        x = self.conv1([x, a])
        # 遍历LayerList
        for conv in self.convs:
            x = conv([x, a])
        x = self.pool([x, i])
        x = self.dense1(x)
        x = self.dropout(x)
        return self.dense2(x)

普通列表中的层不会被Keras视为模型的一部分，其变量不会加入trainable_variables，也会导致梯度追踪失败。

（3）修正DQN调用逻辑

在check_IGDQN中，需要基于模型预测计算目标损失，而非使用固定常数：

def check_IGDQN(designName, inputDir):
    number_of_outputs = 128
    layers = 3
    alpha = 5e-4
    gamma = 0.2
    epsilon = 0.3
    dqn = IGDQN(
            number_of_outputs,
            layers,
            alpha,
            gamma,
            epsilon
    )

    inputs = get_inputs()
    state = inputs
    action = dqn.choose_action(state)
    
    # 模拟目标Q值（实际任务中需用DQN目标网络或Bellman方程计算）
    target_q = dqn.model(state)
    # 针对选中的动作更新目标值，这里为示例简化处理
    target_q = tf.tensor_scatter_nd_update(
        target_q,
        indices=[[0, action]],
        updates=[10.0]  # 实际需根据奖励和下一状态计算
    )
    
    # 调用修正后的update方法
    loss = dqn.update(state, target_q)
    print(f"损失值: {loss.numpy()}")

3. 其他注意事项 #

• choose_action中用model.predict()是eager模式调用，若要兼容tf.function，建议直接用model(state)获取预测值。
• 确保get_inputs返回的输入是TensorFlow张量类型（Spektral的DisjointLoader一般会返回张量，可自行确认）。

内容的提问来源于stack exchange，提问作者mindstorm84