TensorFlow递归模型：如何获取while_loop体内中间结果的梯度

这个问题我之前也碰到过：默认情况下tf.while_loop为了计算效率，不会主动留存循环体内的中间状态张量，所以直接想引用s1、s2这类中间值是行不通的。不过我们可以手动修改循环逻辑，把每个时间步的状态收集起来，这样就能拿到这些中间张量，进而计算梯度了。

核心思路 #

• 用tf.TensorArray（TensorFlow专门用于动态收集张量的容器）来存储每个时间步的中间状态
• 修改循环体，每次迭代时将当前状态写入TensorArray
• 循环结束后，将TensorArray转换为普通张量，提取需要的中间状态
• 用tf.gradient计算最终输出对中间状态的梯度

完整示例代码 #

import tensorflow as tf
tf.InteractiveSession()

sequence = tf.constant([2, 3, 4, 5])

# 初始化TensorArray，支持动态扩容，元素类型和初始状态一致
state_array = tf.TensorArray(dtype=tf.int32, size=0, dynamic_size=True)

def body(t, previous_s, state_array):
    intermediate_s = previous_s * sequence[t]
    # 将当前中间状态写入TensorArray的对应位置
    state_array = state_array.write(t, intermediate_s)
    return t + 1, intermediate_s, state_array

t0, s0 = 0, 1
_, final_s, state_array = tf.while_loop(
    # 循环终止条件：t小于序列长度
    lambda t, _, __: t < sequence.shape[0],
    body,
    (t0, s0, state_array)
)

# 将TensorArray转换为普通张量，得到所有中间状态（s1到s4）
all_intermediate_states = state_array.stack()
s1 = all_intermediate_states[0]
s2 = all_intermediate_states[1]
s3 = all_intermediate_states[2]
s4 = final_s

# 打印验证结果
print('所有中间状态:', all_intermediate_states.eval())  # 输出 [2, 6, 24, 120]
print('最终状态s4:', s4.eval())  # 输出 120

# 计算梯度
ds4_ds1 = tf.gradient(s4, s1)[0]
print('dS4/dS1:', ds4_ds1.eval())  # 输出 60（对应3*4*5）

ds4_ds2 = tf.gradient(s4, s2)[0]
print('dS4/dS2:', ds4_ds2.eval())  # 输出 20（对应4*5）

ds4_ds3 = tf.gradient(s4, s3)[0]
print('dS4/dS3:', ds4_ds3.eval())  # 输出 5

补充说明 #

• 如果你的序列长度是固定的，也可以用tf.concat来收集中间状态，但TensorArray更适配tf.while_loop的动态迭代场景，效率更高
• 如果你原本是用tf.scan实现的递归模型，其实tf.scan本身会返回所有时间步的输出张量，直接从这个输出里提取中间状态即可，不需要额外收集，会更简单

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