如何将 gradientTape() 与不可训练的变量一起使用

Question

假设我有一个 2 层神经网络，我想让第二层不可训练。 所以我启动这些变量

w1 = tf.random.truncated_normal([28*28, 256])
b1 = tf.zeros([256])
w2 = tf.random.truncated_normal([256, 50])
b2 = tf.zeros([10])

并训练他们。

for (x,y) in db:
    x = tf.reshape(x, [-1, 28*28])
    with tf.GradientTape() as tape:
        h1 = x@w1 + tf.broadcast_to(b1, [x.shape[0], 256])
        h1 = tf.nn.relu(h1)
        h2 = h1@w2 + tf.broadcast_to(b2, [x.shape[0], 10])
        out = tf.nn.relu(h2)
        y_onehot = tf.one_hot(y, depth=10)
        loss = tf.square(y_onehot - out)
        loss = tf.reduce_mean(loss)

问题是，随着 GradientTape() 包含所有变量，w2 和 b2 也是可训练的。 如何使它们不可训练？

Answer 1

您假设GradientTape内的所有变量都是可训练的是不正确的。

仅针对您作为第二个参数传递给gradient function 的变量计算梯度：

tape.gradient(tensor, <your trainable variables here>).