在 LSTM tensorflow 2.0 中计算给定时间步的输出相对于输入的导数

Question

我编写了一个示例代码来生成我在项目中面临的真正问题。我在 tensorflow 中使用 LSTM 对一些时间序列数据进行建模。输入维度为(10, 100, 1)，即10个实例，100个时间步长，特征个数为1。输出的形状相同。

我在训练模型后想要实现的是研究每个特定时间步的每个输入对每个输出的影响。换句话说，我想在每个时间步查看哪些输入变量对我的输出影响最大（或者哪个输入对输出的影响最大/可能是大梯度）。下面是这个问题的代码：

tf.keras.backend.clear_session()
tf.random.set_seed(42)

model_input = tf.data.Dataset.from_tensor_slices(np.random.normal(size=(10, 100, 1)))
model_input = model_input.batch(10)
model_output = tf.data.Dataset.from_tensor_slices(np.random.normal(size=(10, 100, 1)))
model_output = model_output.batch(10)

my_dataset = tf.data.Dataset.zip((model_input, model_output))

m_inputs = tf.keras.Input(shape=(None, 1))

lstm_outputs = tf.keras.layers.LSTM(32, return_sequences=True)(m_inputs)
m_outputs = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(1))(lstm_outputs)

my_model = tf.keras.Model(m_inputs, m_outputs, name="my_model")

my_optimizer=tf.keras.optimizers.Adam(learning_rate=0.001)
my_loss_fn = tf.keras.losses.MeanSquaredError()

my_epochs = 3

for epoch in range(my_epochs):

    for step, (x_batch_tr, y_batch_tr) in enumerate(my_dataset):
        x += 1
        # open a gradient tape to record the operations run during the forward pass, which enables autodifferentiation
        with tf.GradientTape() as tape:

            # Run the forward pass of the layer
            logits = my_model(x_batch_tr, training=True)

            # compute the loss value for this mismatch
            loss_value = my_loss_fn(y_batch_tr, logits)

        # use the gradient tape to automatically retrieve the gradients of the trainable variables with respect to the loss.
        grads = tape.gradient(loss_value, my_model.trainable_weights)

        # Run one step of gradient descent by updating the value of the variables to minimize the loss.
        my_optimizer.apply_gradients(zip(grads, my_model.trainable_weights))

        print(f"Step {step}, loss: {loss_value}")


print("\n\nCalculate gradient of ouptuts w.r.t inputs\n\n")

for step, (x_batch_tr, y_batch_tr) in enumerate(my_dataset):
    # open a gradient tape to record the operations run during the forward pass, which enables autodifferentiation
    with tf.GradientTape() as tape:

        tape.watch(x_batch_tr)

        # Run the forward pass of the layer
        logits = my_model(x_batch_tr, training=True)
        #tape.watch(logits[:, 10, :])   # this didn't help
        # compute the loss value for this mismatch
        loss_value = my_loss_fn(y_batch_tr, logits)

    # use the gradient tape to automatically retrieve the gradients of the trainable variables with respect to the loss.
#     grads = tape.gradient(logits, x_batch_tr)   # This works
#     print(grads.numpy().shape)                  # This works
    grads = tape.gradient(logits[:, 10, :], x_batch_tr)
    print(grads)

换句话说，我想关注对我的输出影响最大的输入（在每个特定时间步）。

对我来说 grads = tape.gradient(logits, x_batch_tr) 不会做这项工作，因为这将添加所有输出 w.r.t 每个输入的梯度。

但是，渐变总是无。

非常感谢任何帮助！

Answer 1

您可以使用tf.GradientTape.batch_jacobian 准确获取该信息：

grads = tape.batch_jacobian(logits, x_batch_tr)
print(grads.shape)
# (10, 100, 1, 100, 1)

在这里，grads[i, t1, f1, t2, f2] 为您提供，例如 i，输出特征 f1 在时间 t1 相对于输入特征 f2 在时间 t2 的梯度。如果像您的情况一样，您只有一个功能，您可以说grads[i, t1, 0, t2, 0] 为您提供了t1 相对于t2 的梯度。方便的是，您还可以聚合此结果的不同轴或切片以获得聚合梯度。例如，tf.reduce_sum(grads[:, :, :, :10], axis=3) 将为您提供每个输出时间步相对于前十个输入时间步的梯度。

关于在你的例子中得到None梯度，我认为是因为你在梯度磁带上下文之外进行切片操作，所以梯度跟踪丢失了。

Answer 2

所以解决方案是为我们需要在 tape.grad 中使用的部分 logits 创建一个临时张量，并使用 tape.watch 在磁带上注册该张量

应该是这样的：

for step, (x_batch_tr, y_batch_tr) in enumerate(my_dataset):
    # open a gradient tape to record the operations run during the forward pass, which enables autodifferentiation
    with tf.GradientTape() as tape:

        tape.watch(x_batch_tr)

        # Run the forward pass of the layer
        logits = my_model(x_batch_tr, training=True)
        tensor_logits = tf.constant(logits[:, 10, :])
        tape.watch(tensor_logits)   # this didn't help

        # compute the loss value for this mismatch
        loss_value = my_loss_fn(y_batch_tr, logits)

    # use the gradient tape to automatically retrieve the gradients of the trainable variables with respect to the loss.
    grads = tape.gradient(tensor_logits, x_batch_tr)
    print(grads.numpy())