通过Tensorflow查找并使用损失wrt输入的梯度

Question

我对Tensorflow还是陌生的，尽管我过去使用过Keras，但对此我感到很满意。我正在尝试复制用Pytorch编写的论文（关于ODIN的论文）的结果，我想使用Tensorflow进行相同的操作。基本上，我需要：

1. 获取相对于输入的损耗梯度。
2. 评估图像上的渐变。
3. 将渐变和图像加在一起（因此会干扰原始图像）。
4. 将此新的“图像”提供给网络。

所以我有一个经过MNIST训练的玩具神经网络：

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])
loss='sparse_categorical_crossentropy'
model.compile(optimizer='adam',
              loss=loss,
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test,  y_test, verbose=2)

我在这里读到我可以通过运行以下命令找到相对于输入的渐变。

grads = K.gradients(model.output, model.input)
grads[0]

但是现在我想执行其余的步骤：评估图像上的梯度，添加梯度以扰动图像，然后调用预测函数。所以我有点喜欢

tt = tf.add(model.input, grads[0])
model.predict(tt)

显然这是行不通的，但是我只是无法将自己的头围绕Tensorflow及其语法和工作方式。我知道这一定很愚蠢而且容易做到，但是我不知道应该如何进行。如何有效评估和总结图像的梯度？