TensorFlow损失依赖复杂时梯度计算求助：梯度全零问题

核心问题 #

你代码中使用self.NET2.weights[j].assign(weights[j])给NET2权重赋值的操作，属于变量的副作用修改。TensorFlow的GradientTape无法追踪这种赋值操作与后续NET2前向传播之间的依赖关系，直接切断了从loss到pred_weights的梯度传播路径，最终导致NET1的梯度计算为零。

解决方案：用函数式前向传播替代变量赋值 #

不要修改NET2的变量权重，而是将NET2的前向传播逻辑改为接受权重参数的形式，让GradientTape能完整追踪整个计算图的依赖链。

方案1：抽离NET2的前向计算函数 #

如果NET2是自定义的简单网络，可以直接抽离前向逻辑为独立函数：

def net2_forward(inputs, weights):
    # 按照NET2的实际层结构，用传入的weights完成计算
    # 示例：假设NET2是两层全连接网络
    x = tf.matmul(inputs, weights[0]) + weights[1]
    x = tf.nn.relu(x)
    x = tf.matmul(x, weights[2]) + weights[3]
    return x

修改后的train_step：

def train_step(self, input_weights):
    with tf.GradientTape() as tape:
        pred_weights = self.NET1(input_weights)
        weights = self.transform_weights_from_array(pred_weights)
        # 直接用传入的权重计算NET2输出，不修改原网络变量
        u = net2_forward(SOME_INPUT, weights)
        loss = tf.reduce_sum(tf.math.abs(u))
    
    gradients = tape.gradient(loss, self.NET1.trainable_variables,
                              unconnected_gradients=tf.UnconnectedGradients.ZERO)

方案2：修改NET2的call方法支持外部权重 #

如果NET2是基于tf.keras.Model构建的，可以重写call方法，让它支持传入外部权重：

class NET2(tf.keras.Model):
    def __init__(self):
        super().__init__()
        # 定义层结构（无需初始化权重，或仅作为形状参考）
        self.dense1 = tf.keras.layers.Dense(64, use_bias=True)
        self.dense2 = tf.keras.layers.Dense(10, use_bias=True)
    
    def call(self, inputs, weights=None):
        if weights is not None:
            # 使用传入的权重执行前向计算
            x = tf.matmul(inputs, weights[0]) + weights[1]
            x = tf.nn.relu(x)
            x = tf.matmul(x, weights[2]) + weights[3]
            return x
        # 保留原有的默认前向逻辑（如果需要）
        return super().call(inputs)

对应的train_step修改为：

def train_step(self, input_weights):
    with tf.GradientTape() as tape:
        pred_weights = self.NET1(input_weights)
        weights = self.transform_weights_from_array(pred_weights)
        # 传入外部权重调用NET2
        u = self.NET2(SOME_INPUT, weights=weights)
        loss = tf.reduce_sum(tf.math.abs(u))
    
    gradients = tape.gradient(loss, self.NET1.trainable_variables,
                              unconnected_gradients=tf.UnconnectedGradients.ZERO)

额外注意点 #

• 移除了persistent=True，因为仅需计算一次梯度，该参数会额外占用内存，没必要保留。
• 确保transform_weights_from_array中所有操作都是TensorFlow可微分的（当前实现的tf.reshape是可微分的，没问题）。

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