KerasのModelとLayersを使いながらLambda等一部カスタムなレイヤーと変数も含めて学習してほしい場合の書き方

下記を参考にした

https://www.tensorflow.org/api_docs/python/tf/keras/layers/Lambda

L2 Constrain Softmax Lossを今回実装したが、その時にもどう実装するのかがわかりにくかったのでメモ.

L2 Constrain softmax lossをTensorFlow + Kerasでmnistをとりあえず動かす

本編

Kerasのlayersに含まれていない層で簡単な処理を実現する必要があった. また訓練可能な変数もその層の計算に含まれている.

今回実装したL2正則化してスケールするだけのような処理を考えるにあたって、最初はLambda Layerでできそうかと思った.

alpha = 30
def l2_constrain(x):
    return alpha * K.l2_normalize(x, axis=1)
model.add(layers.Lambda(l2_constrain))

こんな感じで書けば、Lambdaレイヤに関数を登録して、できそうだし、実際に処理できている.しかし、問題はもしalphaを学習させたい場合、どうするかだ.

https://www.tensorflow.org/api_docs/python/tf/keras/layers/Lambda#variables

にも書かれているように、学習させたい対象をレイヤが追えていないと学習時に訓練可能変数として認識してもらえない.

もし訓練ができる変数であることをmodelに伝えるには、カスタムなレイヤーにするしかないのだそう.

なので、下記のような実装になる.

class L2ConstrainLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(L2ConstrainLayer, self).__init__()
        self.alpha = tf.Variable(30.)

    def call(self, inputs):
        return K.l2_normalize(inputs, axis=1) * self.alpha

model.add(L2ConstrainLayer())

この書き方をすればKerasのこれまでの書き方でl2_normalizeも使用することが可能になった.