在扁平参数张量上使用 tf.gradients 或 tf.hessians

Question

假设我想针对某些参数 W（例如，前馈神经网络的权重和偏差）计算标量值函数的 Hessian 矩阵。如果考虑以下代码，实现经过训练以最小化 MSE 损失的二维线性模型：

import numpy as np
import tensorflow as tf

x = tf.placeholder(dtype=tf.float32, shape=[None, 2])  #inputs
t = tf.placeholder(dtype=tf.float32, shape=[None,])  #labels
W = tf.placeholder(np.eye(2), dtype=tf.float32)  #weights

preds = tf.matmul(x, W)  #linear model
loss = tf.reduce_mean(tf.square(preds-t), axis=0) #mse loss

params = tf.trainable_variables() 
hessian = tf.hessians(loss, params)

你会期望 session.run(tf.hessian,feed_dict={}) 到 return 一个 2x2 矩阵（等于 W）。事实证明，因为 params 是一个 2x2 张量，所以输出是一个形状为 [2, 2, 2, 2] 的张量。虽然我可以很容易地重塑张量以获得我想要的矩阵，但当 params 成为不同大小的张量列表时（例如，当模型是一个深度神经网络时），这个操作似乎非常麻烦.

似乎有两种解决方法：

将 params 展平为一维张量，称为 flat_params:
```
flat_params = tf.concat([tf.reshape(p, [-1]) for p in params])
```
所以 tf.hessians(loss, flat_params) 自然地 return 是一个 2x2 矩阵。然而，正如 Why does Tensorflow Reshape tf.reshape() break the flow of gradients? 中针对 tf.gradients 所指出的（但也适用于 tf.hessians），tensorflow 无法在 params 和 params 之间的图中看到符号 link flat_params 和 tf.hessians(loss, flat_params) 将引发错误，因为梯度将被视为 None。
在https://afqueiruga.github.io/tensorflow/2017/12/28/hessian-mnist.html中，代码作者反其道而行之，首先创建平面参数并将其部分重塑为self.params。这个技巧确实有效，并让你得到具有 expected 形状（2x2 矩阵）的 hessian。然而，在我看来，当你有一个复杂的模型时，这会很麻烦，如果你通过内置函数（比如 tf.layers.dense, ..）创建你的模型，则不可能应用。

当 self.params 是任意形状的张量列表时，是否没有从 tf.hessians 获取 Hessian 矩阵（如本例中的 2x2 矩阵）的直接方法？如果不是，您如何自动重塑 tf.hessians 的输出张量？

Answer 1

事实证明（根据 TensorFlow r1.13）如果 len(xs) > 1，则 tf.hessians(ys, xs) returns 张量仅对应于完整的 Hessian 矩阵。本文中的完整故事和解决方案 https://arxiv.org/pdf/1905.05559, and code at https://github.com/gknilsen/pyhessian

在扁平参数张量上使用 tf.gradients 或 tf.hessians

Use tf.gradients or tf.hessians on flattened parameter tensor

python

deep-learning

tensorflow

hessian-matrix