[Keras] モデルのヘシアンと勾配を計算する

# はじめに
この記事ではKerasでヘシアン(ヘッセ行列)を計算する方法について解説します。
ヘシアンは自然勾配を利用した効率的な学習や、モデルの圧縮(Optimal brain damage等)などに利用されています。
計算量の観点からヘシアンがそのまま利用されることは稀ですが、小さいモデルでは現実的な時間で求まるので、今回はその方法について説明していきます。
ソースコードは[こちら](https://github.com/shiita0903/KerasHessian)からどうぞ

# 環境
||version|
|:--|:--|
|Python|3.6.9|
|Keras|2.2.5|
|TensorFlow|1.14.0|

# ヘシアンを計算する対象のモデル
今回ヘシアンを計算する対象は、非常に小さい全結合層のモデルです。

```
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
Dense1 (Dense)               (None, 2)                 4         
_________________________________________________________________
Dense2 (Dense)               (None, 2)                 4         
=================================================================
Total params: 8
Trainable params: 8
Non-trainable params: 0
_________________________________________________________________
```

簡単のためにバイアスは利用していません。活性化関数は`ReLU`で、最終層はクラス分類問題を想定して、`Softmax`を利用しています。
今回は学習をしないので、`Model`の`set_weights()`メソッドで重みの値を適当にセットしておきます。

```python
def build_model() -> Model:
    model = Sequential([
        Dense(2, use_bias=False, activation="relu", name="Dense1", input_shape=(2,)),
        Dense(2, use_bias=False, activation="softmax", name="Dense2")
    ])
    model.summary()
    model.set_weights([
        np.array([[1, 2], [3, 4]]),  # Dense1の重み
        np.array([[6, 5], [7, 8]]),  # Dense2の重み
    ])
    return model
```

# ヘシアンの計算
必要な`import`は以下の通りです。

```python
import keras.backend as K
import numpy as np
import tensorflow as tf
from keras import Model, Sequential
from keras.layers import Dense
```

ヘシアンを計算するためにまず、`Loss`を返すテンソルを作成します。
` Model`に設定された`Optimizer`から`Loss`を返すテンソルを取得する方法がわからなかったので(そもそも出来ないかも)、Kerasのバックエンド関数を活用して自分で作っていきます。

```python
model = build_model()
y_true = K.placeholder((None, 2,))  # one-hotなラベルを入れるPlaceholder
loss = K.categorical_crossentropy(y_true, model.output)  # model.outputはSoftmax後の値
```

Kerasのバックエンド関数ではヘシアンを計算出来ないので、やむを得ずTensorFlowの関数を使います。
大多数の人はバックエンドにTensorFlowを利用していると思われるので、問題は無いはず。
また、`tf.hessians()`は4次元のテンソルを返すので、見やすいように2次元テンソルに変形します。

```python
hessian = tf.hessians(loss, model.get_layer("Dense1").kernel)[0]
s = hessian.shape
hessian = K.reshape(hessian, [s[0] * s[1], s[2] * s[3]])  # 4次元テンソルを2次元に整形
```

最後に`eval()`メソッドに`feed_dict`としてモデルのインプットとラベルを渡してあげて、ヘシアンの計算をします。

```python
inputs = np.array([[1, 2]])
labels = np.array([[0, 1]])

with K.get_session():
    print(hessian.eval({model.input: inputs, y_true: labels}))
```

```plaintext:結果
[[ 0.04517667 -0.04517668  0.09035334 -0.09035337]
 [-0.04517668  0.04517663 -0.09035337  0.09035325]
 [ 0.09035334 -0.09035337  0.18070668 -0.18070674]
 [-0.09035337  0.09035325 -0.18070674  0.1807065 ]]
```

# 勾配の計算
同様の手順でモデルの勾配も計算できます。

```python
gradient = K.gradients(loss, model.get_layer("Dense1").kernel)[0]

with K.get_session():
    print(gradient.eval({model.input: inputs, y_true: labels}))
```

```plaintext:結果
[[ 0.0474259  -0.04742587]
 [ 0.09485179 -0.09485173]]
```

# おわりに
コードを見れば分かることですが、かなりシンプルに計算が出来ました。
ですが、Kerasのバックエンド関数が絡んでくる様になると、GitHubのコードを読んで内部を理解したりとかなり骨が折れます。
`placeholder`や`feed_dict`を利用したテンソルの評価や、`Model`からレイヤを取得する方法、モデルの重みの取得方法や設定方法など、それなりに参考になるかと思います。