tf.GradientTape结合外积操作返回None梯度的问题咨询

核心错误点 #

• 梯度链路被人为断开：你在计算u_pred时用tf.Variable()对运算结果做了包裹。tf.Variable是TensorFlow专门用于定义模型可训练参数的接口，将中间计算结果套入该接口后，会生成一个独立于原计算图的新变量，梯度无法从u_pred回传到模型输出y_pred，最终导致模型参数梯度全为None。
• 逐样本循环的实现方式存在冗余：用列表推导加for循环逐样本计算外积的实现不仅运行效率低，还可能隐式破坏计算图的依赖关系，即便去掉tf.Variable也容易出现性能或梯度异常。

修正后的代码实现 #

直接使用批量向量化运算实现需求，不需要逐样本循环，也不需要套tf.Variable：

with tf.GradientTape() as tape:
    y_pred = model(X_batch, training=True)
    # 拆分f、g并扩展维度用于批量外积计算
    f = y_pred[:, :nX, tf.newaxis]
    g = y_pred[:, nX:, tf.newaxis]
    # 批量计算外积，维度为[batch_size, nX, nT]
    outer_product = f @ tf.transpose(g, perm=(0, 2, 1))
    # 按Fortran序展平，和原numpy实现的flatten('F')效果完全一致
    u_pred = tf.reshape(tf.transpose(outer_product, perm=(0, 2, 1)), (-1, nX * nT))
    
    main_loss = tf.reduce_mean(loss_fn(y_batch, u_pred))
    loss = tf.add_n([main_loss] + model.losses)
gradients = tape.gradient(loss, model.trainable_variables)

该实现和你原numpy版本的输出数值完全一致，且全程保留计算图的梯度传播链路，梯度可以正常回传到模型参数。

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