Ahogrammer

Deep Dive Into NLP, ML and Cloud

Kerasでカスタムレイヤーを作成する方法

Kerasでは様々なレイヤーが事前定義されており、それらをレゴブロックのように組み合わせてモデルを作成していきます。 たとえば、EmbeddingやConvolution, LSTMといったレイヤーが事前定義されています。 通常は、これらの事前定義された便利なレイヤーを使ってモデルを作成します。

しかし、事前定義されたレイヤーだけでは必要な機能を実装できない時があります。 たとえば、最新の論文を実装しようとしたときに、その中で使われている機能がKerasに存在しないということはよくあります。

必要なレイヤーがKerasに存在しない場合、自分でカスタムレイヤーを実装する必要があります。 Kerasは自分でカスタムレイヤーを実装するための機能を提供してくれています。 基本的に重みを持たないレイヤーに関してはLambdaで定義し、学習可能な重みを持たせたい場合はカスタムレイヤーの実装を検討することになります。

本記事では、Kerasでのカスタムレイヤーの作成方法を紹介します。

カスタムレイヤーの作り方

カスタムレイヤーを作成するためには、keras.engine.topologyモジュールのLayerクラスを継承したクラスを作る必要があります。 具体的には、Layerで定義されているメソッドをオーバーライドして、必要なレイヤーを作成します。

その際に重要なメソッドは以下の4つです:

  • __init__
  • build
  • call
  • compute_output_shape

以下では各メソッドの役割について説明していきます。

__init__()

初期化メソッドでは、レイヤーを初期化します。 具体的には、必要な機能を実装する際に必要な情報を渡しておきます。 たとえば、重みの初期化方法とか内部のユニット数とかそんな情報を渡します。

build(self, input_shape)

buildでは学習可能な重みを定義します。 具体的には、self.add_weightsメソッドにshapeや初期化方法、正規化方法を渡して重みを定義します。 ここで使用する初期化方法や正規化方法は__init__()メソッドに渡しておいたものを使用します。

また、buildの最後には必ず、self.build = Trueと設定しなければなりません。 self.buildに直接設定してもいいのですが、super([Layer], self).build()を呼び出してもできます。

call(self, x, mask=none)

callメソッドにはForwardの計算で行いたいことを実装します。 callメソッドの中では、__init__で設定した情報、buildで定義した重み、そしてレイヤーへの入力を使って、自分の行いたい計算をします。 カスタムレイヤーのメソッドの中でも最も重要なメソッドの一つです。

compute_output_shape(self, input_shape)

このメソッドは名前の通り出力のshapeを計算します。 作成したレイヤーの内部で入力のshapeを変更する場合には,ここでshape変換のロジックを指定します。 変更しない場合は定義する必要はありません。入力のshapeがそのまま返されます。

カスタムレイヤの保存と読み込み

Kerasではモデルを保存する際に、レイヤーの情報をjsonに変換します。 その際、レイヤーのget_configメソッドが重要な役割を果たします。 get_configメソッド内に、どの情報をjson化するのか記述していきます。 このjsonには、カスタムレイヤーを復元するのに必要な情報を含める必要があります。

カスタムレイヤーの読み込みについては公式のFAQに方法が示されています。
Handling custom layers (or other custom objects) in saved models

簡単に述べておくと、load_model関数の引数であるcustom_objectsに定義したカスタムレイヤーを渡しましょうという話です。 以下のような感じで渡すと、保存したモデルを正常に読み込むことができます。

from keras.models import load_model
# AttentionLayer というカスタムレイヤーを定義したと仮定
model = load_model('my_model.h5', custom_objects={'AttentionLayer': AttentionLayer})

各メソッドの呼ばれる順番

ここでは各メソッドが呼ばれる順番について書いておきます。 呼ばれる順番について知らなくてもレイヤーを定義することはできますが、知っているとデバッグするときに役立つかもしれません。

まず、レイヤーをインスタンス化する際に __init__ が呼ばれます。これは問題ないでしょう。 レイヤーをインスタンス化したら、レイヤーに入力を与えて呼び出します。 一応書いておくと、これは以下のようなイメージです:

x = Input(batch_shape=(None, maxlen), dtype="int32")
embedding = Embedding(vocab_size, dim)(x)

上記のEmbeddingレイヤーのように、レイヤーを呼び出した際に実行されるのが、__call__メソッドです。 __call__ の中では以下の順序でメソッドが呼ばれます:

  • build
  • call
  • compute_mask
  • compute_output_shape

ちなみに、compute_maskではマスキングの計算を行います。

実際のコード

カスタムレイヤーの公式Exampleとして、以下のファイルが挙げられます。
keras/antirectifier.py at master · fchollet/keras · GitHub

上記のExampleで基本的なところを理解したら、Kerasに事前定義されているレイヤーを見るのが一番勉強になるかなと感じています。

ちなみに、私もカスタムレイヤーを定義してみました。Keras用のCRF層です。
keras-crf-layer

参考資料

カスタムレイヤーの作り方に関する公式ドキュメント
Writing your own Keras layers

カスタムレイヤーの公式Example
keras/antirectifier.py at master · fchollet/keras · GitHub

カスタムレイヤーの作成方法のチュートリアル
For beginners; Writing a custom Keras layer

カスタムレイヤーでシリアライズを行う話
kerasでカスタムレイヤーのシリアライズを行う

カスタムレイヤーの読み込み方法
Handling custom layers (or other custom objects) in saved models