tensorflow 不更新不可训练变量的全变量张流反向传播

qnakjoqk  于 2023-01-21  发布在  其他
关注(0)|答案(1)|浏览(127)

我想计算可训练和不可训练变量的梯度。
并且仅更新可训练参数。
首先,我按如下方式实现它

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

gradients = g.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))

non_train_gradients = g.gradient(loss, model.non_trainable_variables)

然而,上面的代码做了两次反向传播来计算梯度。
我想同时估计可训练和不可训练变量的梯度,
而仅更新可训练参数。
我该怎么做呢?

5t7ly7z5

5t7ly7z51#

我们可以利用梯度只是一个列表并且按照与我们放入的变量相同的顺序返回的事实:

n_trainable = len(model.trainable_variables)
    
gradients = g.gradient(
    loss, model.trainable_variables + model.non_trainable_variables
)

trainable_gradients = gradients[:n_trainable]
non_trainable_gradients = gradients[n_trainable:]

optimizer.apply_gradients(
    zip(trainable_gradients, model.trainable_variables)
)

也就是说,我们只是把所有不可训练的变量放在最后,然后在那一点上分割梯度。

相关问题