自定义TensorFlow GCN层实现模型无法学习问题排查求助

核心错误点 #

• 可训练变量完全没有被跟踪和更新
  你定义的GCNLayer和MyModel都是普通Python类，没有继承TensorFlow的tf.keras.layers.Layer和tf.keras.Model基类，且每次调用层的__call__方法时都会重新实例化Dense层，参数不会被持久化保存，也不会被梯度磁带自动识别为可训练变量。你可以打印训练时的variables变量，大概率是空列表，所以整个训练过程参数根本没有更新。
• 交叉熵损失参数传递错误
  tf.nn.softmax_cross_entropy_with_logits要求必须显式传递关键字参数，你按位置传递参数会导致labels和logits传反，损失计算完全错误，正确写法是tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=logits)。
• 导入语法错误
  代码中import tensorflow as 后面漏掉了tf，属于基础语法错误。
• GCN层逻辑的潜在问题（非核心但建议调整）
  标准GCN的计算顺序是先聚合邻域特征再做线性变换，虽然从矩阵乘法结合律来看你当前先变换后聚合的写法数学上等价，但推荐遵循标准实现逻辑避免后续扩展出错。

修复后的代码参考 #

!pip install numpy
!pip install tensorflow
!pip install spektral

import numpy as np
import tensorflow as tf
import spektral

def masked_cross_entropy_loss(labels, logits, mask):
  # 修复：使用关键字参数传递
  loss = tf.nn.softmax_cross_entropy_with_logits(labels=labels, logits=logits)
  mask = tf.cast(mask, dtype=tf.float32)
  mask /= tf.reduce_mean(mask)
  loss *= mask
  return tf.reduce_mean(loss)

def masked_accuracy(labels, logits, mask):
  accuracy_array = tf.equal(tf.argmax(logits, axis=1), tf.argmax(labels, axis=1))
  accuracy_array = tf.cast(accuracy_array, dtype=tf.float32)
  mask = tf.cast(mask, dtype=tf.float32)
  mask /= tf.reduce_mean(mask)
  accuracy_array *= mask
  return tf.reduce_mean(accuracy_array) 

# 修复：继承keras层基类
class GCNLayer(tf.keras.layers.Layer):
  def __init__(self, A, units, activation='relu', **kwargs):
    super().__init__(**kwargs)
    self.A = A
    self.units = units
    self.activation = tf.keras.activations.get(activation)
    # 修复：在初始化阶段定义可训练层，不要每次call时新建
    self.dense = tf.keras.layers.Dense(units)
  
  def call(self, features):
    # 可调整为标准顺序：先聚合再变换
    features = self.dense(features)
    features = self.A @ features
    if self.activation is not None:
      features = self.activation(features)
    return features

# 修复：继承keras模型基类
class MyModel(tf.keras.Model):
  def __init__(self, A, node_features, node_labels, train_mask, val_mask, test_mask, **kwargs):
    super().__init__(**kwargs)
    self.A = A
    self.node_features = node_features
    self.node_labels = node_labels
    self.train_mask = train_mask
    self.val_mask = val_mask
    self.test_mask = test_mask
    
    self.gcn_layer1 = GCNLayer(self.A, units=32, activation='relu')
    self.gcn_layer2 = GCNLayer(self.A, units=7, activation=None)

  def call(self, inputs=None):
    hidden_out = self.gcn_layer1(self.node_features)
    output = self.gcn_layer2(hidden_out)
    return output

  def train(self, num_epochs=200, lr=0.01):
    optimizer = tf.keras.optimizers.Adam(lr)
    best_val_acc = 0.0
    for e in range(num_epochs):
      with tf.GradientTape() as t:
        logits = self()
        train_loss = masked_cross_entropy_loss(self.node_labels, logits, self.train_mask)
      
      # 修复：直接获取模型的可训练变量
      grads = t.gradient(train_loss, self.trainable_variables)
      optimizer.apply_gradients(zip(grads, self.trainable_variables))
      logits = self()
      train_acc = masked_accuracy(self.node_labels, logits, self.train_mask)
      val_acc = masked_accuracy(self.node_labels, logits, self.val_mask)
      if val_acc > best_val_acc:
        best_val_acc = val_acc
        print(f'epoch={e},Training Loss:{train_loss.numpy()},Training Accuracy:{train_acc.numpy()}, Validation Accuracy:{val_acc.numpy()}')

# 假设A、node_features等变量已提前定义
model = MyModel(A, node_features, node_labels, train_mask, val_mask, test_mask)
model.train(num_epochs=200, lr=0.01)

额外验证建议 #

• 确认邻接矩阵的归一化方式是标准的D^(-0.5) * (A + I) * D^(-0.5)，添加自环避免节点自身特征丢失。
• 训练初期可以打印梯度值，确认梯度不为None，参数确实在更新。

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