使用Keras获取模型输出w.r.t权重的梯度

import tensorflow as tf
from tensorflow import keras

model = keras.Sequential([
    keras.layers.Dense(units=64, activation='relu', input_shape=(784,)),
    keras.layers.Dense(units=10, activation='softmax')
])

loss_fn = keras.losses.SparseCategoricalCrossentropy()

optimizer = keras.optimizers.SGD(learning_rate=0.001)

@tf.function
def get_gradient(inputs, targets):
    with tf.GradientTape() as tape:
        predictions = model(inputs)
        loss_value = loss_fn(targets, predictions)
    return tape.gradient(loss_value, model.trainable_weights)

inputs = ...
targets = ...
gradients = get_gradient(inputs, targets)