如何在tensorflow 2中获得损失梯度wrt内层输出？

Question

我想在训练期间获得模型损失函数相对于特定层输出的梯度。接下来我想用它做的是在下一个学习时期使用该梯度的值来修改图层中的某些内容。那么如何获得那个梯度呢？

这是一个最小的例子。 MinimalRNNCell 代码是从 TensorFlow 的网站复制的，玩具数据仅用于重现该行为。

import tensorflow as tf 
from tensorflow.keras.layers import RNN, SimpleRNNCell, SimpleRNN, Layer, Dense, AbstractRNNCell
from tensorflow.keras import Model
import numpy as np
import tensorflow.keras.backend as K


class MinimalRNNCell(AbstractRNNCell):

    def __init__(self, units, **kwargs):
      self.units = units
      super(MinimalRNNCell, self).__init__(**kwargs)

    @property
    def state_size(self):
      return self.units

    def build(self, input_shape):
      self.kernel = self.add_weight(shape=(input_shape[-1], self.units),
                                    initializer='uniform',
                                    name='kernel')
      self.recurrent_kernel = self.add_weight(
          shape=(self.units, self.units),
          initializer='uniform',
          name='recurrent_kernel')
      self.built = True

    def call(self, inputs, states):
      prev_output = states[0]
      h = K.dot(inputs, self.kernel)
      output = h + K.dot(prev_output, self.recurrent_kernel)
      return output, output


class MyModel(Model):
    def __init__(self, size):
        super(MyModel, self).__init__()
        self.minimalrnn=RNN(MinimalRNNCell(size), name='minimalrnn')
        self.out=Dense(4)

    def call(self, inputs):
        out=self.minimalrnn(inputs)
        out=self.out(out)
        return out


x=np.array([[[3.],[0.],[1.],[2.],[3.]],[[3.],[0.],[1.],[2.],[3.]]])
y=np.array([[[0.],[1.],[2.],[3.]],[[0.],[1.],[2.],[3.]]])

model=MyModel(2)
model.compile(optimizer='sgd', loss='mse')
model.fit(x,y,epochs=10, batch_size=1, validation_split=0.2)

现在我想获得 MyModel 的 minimalrnn 层输出的梯度（在每批数据之后）。

如何做到这一点？我想我可以尝试使用 GradientTape 观看 model.get_layer('minimalrnn').output，但我需要更多学习资源或示例。

编辑

我在 Tiago Martins Peres 提供的代码中使用了 GradientTape，但我特别想要获得梯度 wrt 层输出，但我仍然无法实现。

现在在 class 定义之后，我的代码如下所示：


x=np.array([[[3.],[0.],[1.],[2.],[3.]],[[3.],[0.],[1.],[2.],[3.]]])
y=np.array([[0., 1., 2., 3.],[0., 1., 2., 3.]])

model=MyModel(2)

#inputs = tf.keras.Input(shape=(2,5,1))
#model.call(x)

def gradients(model, inputs, targets):
    with tf.GradientTape() as tape:
        tape.watch(model.get_layer('minimalrnn').output)
        loss_value = loss_fn(model, inputs, targets)
    return tape.gradient(loss_value, model.trainable_variables)

def loss_fn(model, inputs, targets):
    error = model(inputs) - targets
    return tf.reduce_mean(tf.square(error))

optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
print("Initial loss: {:.3f}".format(loss_fn(model, x, y)))
for i in range(10):
    grads = gradients(model, x, y)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    print("Loss at step {:03d}: {:.3f}".format(i, loss_fn(model, x, y)))
print("Final loss: {:.3f}".format(loss_fn(model, x, y)))

如您所见，我在渐变函数定义中添加了 tape.watch，因为我想观察图层输出。但是我收到错误：

Traceback (most recent call last):
  File "/home/.../test2.py", line 73, in <module>
    grads = gradients(model, x, y)
  File "/home/.../test2.py", line 58, in gradients
    print(model.get_layer('minimalrnn').output)
  File "/home/.../.venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 1553, in output
    raise AttributeError('Layer ' + self.name + ' has no inbound nodes.')
AttributeError: Layer minimalrnn has no inbound nodes.

我还尝试在具有指定大小（注释行）的输入上调用模型，根据对此的回答：。它没有帮助。在模型的初始化函数中指定输入形状，如下所示，也无济于事 - 仍然是同样的错误。

self.minimalrnn=RNN(MinimalRNNCell(size), name='minimalrnn', input_shape=(2,5,1))

Answer 1

是的，您可以使用 GradientTape。 tf.GradientTape 的目的是记录自动微分操作或计算操作或计算相对于其输入变量的梯度。

根据What's New in TensorFlow 2.0，首先使用tf.GradientTape实现模型的简单训练，在tf.GradentTape上下文管理器中调用输入张量的正向传递，然后计算损失函数。 这确保所有的计算都将记录在梯度磁带上。

然后，计算模型中所有可训练变量的梯度。计算出梯度后，可以执行任何所需的梯度裁剪、归一化或转换，然后再将它们传递给优化器以将它们应用于模型变量。看看下面的例子：

NUM_EXAMPLES = 2000

input_x = tf.random.normal([NUM_EXAMPLES])
noise = tf.random.normal([NUM_EXAMPLES])
input_y = input_x * 5 + 2 + noise

def loss_fn(model, inputs, targets):
  error = model(inputs) - targets
  return tf.reduce_mean(tf.square(error))

def gradients(model, inputs, targets):
  with tf.GradientTape() as tape:
    loss_value = loss_fn(model, inputs, targets)
  return tape.gradient(loss_value, model.trainable_variables)

model = tf.keras.Sequential(tf.keras.layers.Dense(1))
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
print("Initial loss: {:.3f}".format(loss_fn(model, input_x, input_y)))
for i in range(500):
  grads = gradients(model, input_x, input_y)
  optimizer.apply_gradients(zip(grads, model.trainable_variables))
  if i % 20 == 0:
    print("Loss at step {:03d}: {:.3f}".format(i, loss_fn(model, input_x, input_y)))
print("Final loss: {:.3f}".format(loss(model, input_x, input_y)))
print("W = {}, B = {}".format(*model.trainable_variables))

Answer 2

好的，我终于找到的一个答案隐藏在这里：。我什至可以使用子类模型。

此外，AttributeError 的问题很奇怪，因为当我使用 Sequential 而不是子类化 Model 时，AttributeError 神奇地消失了，也许它与这个问题有关 https://github.com/tensorflow/tensorflow/issues/34834？

不过，我想知道为什么我不能将层的输出作为第二个参数传递给 tape.gradient。

如何在tensorflow 2中获得损失梯度wrt内层输出？

How to get loss gradient wrt internal layer output in tensorflow 2?

python

machine-learning

tensorflow

tf.keras

tensorflow2.0