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

Question

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

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

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 每个输入。

然而，渐变总是None。

非常感谢任何帮助！

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] 相对于 t2 给你 t1 的梯度。方便的是，您还可以聚合此结果的不同轴或切片以获得聚合梯度。例如，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())

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

Calculating the derivates of the output with respect to input for a give time step in LSTM tensorflow2.0

python

tensorflow

lstm

gradient-descent