TensorFlow中可训练与不可训练变量拼接后的训练属性问题

结论 #

C的前64列可训练，后32列不可训练。

详细解释 #

在TensorFlow中，tf.concat([A, B], axis=1)操作并不会生成一个全新的独立变量，而是创建一个引用原始变量A和B的张量。损失优化过程中，TensorFlow只会为标记为可训练的变量计算梯度并更新其值：

• C的前64列直接对应可训练变量A，反向传播时梯度会回流到A，因此这部分的值会在训练过程中被优化更新。
• C的后32列来自不可训练变量B，TensorFlow不会为B计算梯度，所以这部分的值在训练时保持不变。

你可以通过以下代码直观验证这个逻辑：

import tensorflow as tf

# 初始化变量
A = tf.Variable(tf.random.normal([123, 64]), trainable=True)
B = tf.Variable(tf.random.normal([123, 32]), trainable=False)
C = tf.concat([A, B], axis=1)

# 定义损失和优化器
loss_fn = lambda: tf.reduce_mean(C ** 2)
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)

# 查看梯度情况
with tf.GradientTape() as tape:
    loss_val = loss_fn()
grads = tape.gradient(loss_val, [A, B])

print("A的梯度是否存在:", grads[0] is not None)  # 输出True
print("B的梯度是否存在:", grads[1] is not None)  # 输出False

运行结果显示，只有A会被计算梯度并更新，B则不会，这直接反映到C的对应列上。

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