如何在TensorFlow中同时计算两类变量梯度仅更新可训练参数

原来的代码会执行两次反向传播，根源是连续调用了两次g.gradient()，每次调用都会触发一轮反向计算。要实现一次计算两类变量的梯度，只需要把可训练和不可训练变量合并成一个列表，一次性传给gradient()方法即可，具体代码如下：

with tf.GradientTape() as g:
    preds = model(data)
    loss = criterion(labels, preds)

# 合并可训练与不可训练变量，一次性计算所有梯度
all_variables = model.trainable_variables + model.non_trainable_variables
all_gradients = g.gradient(loss, all_variables)

# 拆分梯度：前半部分对应可训练变量的梯度，后半部分对应不可训练变量的梯度
trainable_gradients = all_gradients[:len(model.trainable_variables)]
non_train_gradients = all_gradients[len(model.trainable_variables):]

# 仅对可训练变量应用梯度更新
optimizer.apply_gradients(zip(trainable_gradients, model.trainable_variables))

关键说明： #

• 不需要再设置persistent=True，因为现在只调用一次g.gradient()，tape用完后会自动释放，更节省内存。
• 一次性计算所有梯度只会触发一次反向传播，大幅降低计算开销。
• 拆分梯度时，利用两类变量的拼接顺序对应梯度列表的顺序，直接通过索引拆分即可。

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