python - 如何在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 层输出，但我仍然无法实现。

现在在类定义之后我的代码如下所示:

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.

我还尝试在具有指定大小(注释行)的输入上调用模型，根据对此的回答:Accessing layer's input/output using Tensorflow 2.0 Model Sub-classing .它没有帮助。在模型的初始化函数中指定输入形状，如下所示，也无济于事 - 仍然是同样的错误。

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

Answer 1

好吧，我终于找到了一个答案隐藏在这里:https://stackoverflow.com/a/56567364/4750170 .我什至可以使用子类模型。

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

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