如何在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 的 minimumrnn 層的輸出梯度（在每批數據之后）。

這該怎么做？ 我想我可以嘗試使用 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.

根據對此的回答，我還嘗試在 Input 上調用具有指定大小（注釋行）的模型： Accessing layer's input/output using Tensorflow 2.0 Model Sub-classing 。 它沒有幫助。 在模型的 init 函數中指定輸入形狀，如下所示，也無濟於事 - 仍然是同樣的錯誤。

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

Answer 1

是的，您可以使用GradientTape 。 tf.GradientTape的目的是記錄用於自動微分或計算操作或計算相對於其輸入變量的梯度的操作。

根據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

好的，我最終找到的一個答案隱藏在這里： https : //stackoverflow.com/a/56567364/4750170 。 我什至可以使用子類模型。

另外 AttributeError 的問題很奇怪，因為當我使用 Sequential 而不是子類化 Model 時，AttributeError 神奇地消失了，也許它與這個問題有關https://github.com/tensorflow/tensorflow/issues/34834 ？

不過，我想知道為什么我不能將層的輸出作為第二個參數傳遞給tape.gradient。