TensorFlow获取中间层输出计算梯度返回None的实现正确性咨询

错误原因 #

你当前的实现方式不正确，梯度返回None的核心原因是计算损失的前向路径和你生成中间层输出的前向路径相互独立，二者没有计算图依赖，梯度带无法跟踪损失和你存入logits的张量之间的关联：

• 你手动逐层调用生成logits里的三个张量，后续又单独调用mnist_model(images, training=True)得到最终输出output，两次调用是完全独立的前向过程，二者的计算节点没有关联
• 直接访问layer.output属性仅适用于静态构建的计算图场景（比如Functional API定义的模型），eager执行下动态生成的张量和当前输入的images没有绑定关系，也会导致梯度断连

正确实现方式 #

方法1：统一前向路径（适配所有模型类型） #

全程用同一条前向路径计算中间输出和最终输出，保证梯度流连贯：

with tf.GradientTape(persistent = True) as tape:
    layer0_out = mnist_model.layers[0](images, training=True)
    layer1_out = mnist_model.layers[1](layer0_out, training=True)
    layer2_out = mnist_model.layers[2](layer1_out, training=True)
    # 如果模型还有后续层，接着往下计算最终输出即可，不要重新调用mnist_model(images)
    output = mnist_model.layers[3](layer2_out, training=True) 
    logits = [layer0_out, layer1_out, layer2_out]
    loss_value = loss_object(labels, output)
    loss_history.append(loss_value.numpy().mean())

grads0 = tape.gradient(loss_value, logits[0])
grads1 = tape.gradient(loss_value, logits[1])
grads2 = tape.gradient(loss_value, logits[2])

方法2：显式监听中间层输出（仅适配Functional/Sequential静态构建的模型） #

不需要拆层计算，用tape.watch显式让梯度带跟踪中间层的输出即可：

with tf.GradientTape(persistent = True) as tape:
    # 显式声明要监听的中间层输出
    tape.watch([mnist_model.layers[0].output, mnist_model.layers[1].output, mnist_model.layers[2].output])
    output = mnist_model(images, training=True)
    loss_value = loss_object(labels, output)
    loss_history.append(loss_value.numpy().mean())

grads0 = tape.gradient(loss_value, mnist_model.layers[0].output)
grads1 = tape.gradient(loss_value, mnist_model.layers[1].output)
grads2 = tape.gradient(loss_value, mnist_model.layers[2].output)

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