Tensorflow：針對每個樣本計算Hessian

Question

我有一個張量X ，大小為M x D。 我們可以將X每一行解釋為訓練樣本，將每一列解釋為特征。

X用於計算大小為M x 1的張量u （換句話說， u取決於計算圖中的X ）。 我們可以將其解釋為預測的載體； 每個樣品一個。 特別地，僅使用X的第m行來計算u第m行。

現在，如果我運行tensor.gradients(u, X)[0] ，我將獲得一個M x D張量，該張量對應於u相對於X的u的“每個樣本”梯度。

我如何才能類似地計算“每個樣本”的黑森州張量？ （即M x D x D數量）

---

附錄 ：以下彼得的答案是正確的。 我還發現了使用堆棧和堆棧的另一種方法（使用彼得的表示法）：

hess2 = tf.stack([
    tf.gradients( tmp, a )[ 0 ]
    for tmp in tf.unstack( grad, num=5, axis=1 )
], axis = 2)

在彼得的示例中， D = 5是要素數量。 我懷疑（但我沒有檢查過）對於M大來說，上述速度更快，因為它跳過了彼得答案中提到的零條目。

Answer 1

tf.hessians()是計算的尺寸的設置YS和xs reagardless Hessian矩陣。 由於您得到的結果尺寸為M x D，而xs的尺寸為M x D，因此結果將為尺寸M x D x M x D。 但是，由於每個示例的輸出彼此獨立，因此大多數Hessian將為零，即，三維中只有一個分片將具有任何值。 因此，要獲得所需的結果，您應該在兩個M維中采用對角線，或者更容易，您應該像這樣簡單地求和並消除第三個維：

hess2 = tf.reduce_sum( hess, axis = 2 )

示例代碼（經過測試）：

import tensorflow as tf

a = tf.constant( [ [ 1.0, 1, 1, 1, 1 ], [ 2, 2, 2, 2, 2 ], [ 3, 3, 3, 3, 3 ] ] )
b = tf.constant( [ [ 1.0 ], [ 2 ], [ 3 ], [ 4 ], [ 5 ] ] )
c = tf.matmul( a, b )
c_sq = tf.square( c )

grad = tf.gradients( c_sq, a )[ 0 ]

hess = tf.hessians( c_sq, a )[ 0 ]
hess2 = tf.reduce_sum( hess, axis = 2 )


with tf.Session() as sess:
    res = sess.run( [ c_sq, grad, hess2 ] )

    for v in res:
        print( v.shape )
        print( v )
        print( "=======================")

將輸出：

（3，1）
[[225.]
[900.]
[2025.]
=======================
（3、5）
[[30. 60. 90. 120. 150.]
[60. 120. 180. 240. 300.]
[90. 180. 270. 360. 450.]]
=======================
（3、5、5）
[[[2. 4. 6. 8. 8. 10.]
[4. 8. 12. 16. 20.]
[6. 12. 18. 24. 30.]
[8. 16. 24. 32. 40.]
[10。 20. 30. 40. 50.]]

[[2. 4. 6. 8. 8. 10.]
[4. 8. 12. 16. 20.]
[6. 12. 18. 24. 30.]
[8. 16. 24. 32. 40.]
[10。 20. 30. 40. 50.]]

[[2. 4. 6. 8. 8. 10.]
[4. 8. 12. 16. 20.]
[6. 12. 18. 24. 30.]
[8. 16. 24. 32. 40.]
[10。 20. 30. 40. 50.]]]
=======================