kerasのモデルの定義方法メモ

最近、kerasで遊ぶ機会が増えてきた。モデルの作り方が複数あってちょっと紛らわしいのでメモがてらに整理する。
なお、自分はtensorflowの上のkerasを利用している。

モデルの定義方法の種類と制限

kerasには主に3種類のモデルの定義方法がある。

Sequenceを利用するもの
Functional APIを利用するもの
Modelクラスを継承するもの

Modelを継承する場合の注意点

Modelを継承する場合には以下の注意が必要である。

model.saveが動作しない、save_weightsは動作するのでそちらを使う

eager_executionと組み合わせるときの注意点

eager_executionを利用するときに以下の制限がある。

optimizerがkeras側のものは動作しないのでtensorflow側のoptimizerを利用する
tensorboardのcallbackが上手く動作しないので自前でやる必要がある

Sequenceを利用するモデル定義

簡単なものだと以下の様なSequenceを使った定義ができる。
とりあえずの動作の確認には便利だが、分岐するネットワークなどの定義は出来ない。

from keras.models import Sequential
from keras.layers import Dense, Activation

model = Sequential([
    Dense(32, input_shape=(784,)),
    Activation('relu'),
    Dense(10),
    Activation('softmax'),
])

Functional APIを利用するモデル定義

Functional APIを利用して以下のようにも定義できる。

from keras.layers import Input, Dense
from keras.models import Model

# This returns a tensor
inputs = Input(shape=(784,))

# a layer instance is callable on a tensor, and returns a tensor
x = Dense(64, activation='relu')(inputs)
x = Dense(64, activation='relu')(x)
predictions = Dense(10, activation='softmax')(x)

# This creates a model that includes
# the Input layer and three Dense layers
model = Model(inputs=inputs, outputs=predictions)
model.compile(optimizer='rmsprop',
              loss='categorical_crossentropy',
              metrics=['accuracy'])
model.fit(data, labels)  # starts training

Modelクラスを継承する定義

Modelクラスを継承して独自のモデルクラスを作成することもできる。
eager_executionと組み合わせると強力。

keras開発者のFrançois Chollet さんの以下のtweetのように、簡潔にRNNなどの記述ができて、デバッグなどもcallに適当にbreak入れたり、printしてデバッグすることも出来ます。

As of TensorFlow 1.7, with eager execution, you can use a Chainer-like API in tf.keras to create models with dynamic control flow. This should be especially useful for NLP researchers working with recursive models! pic.twitter.com/ZfAs3u0HJ0
— François Chollet (@fchollet) March 31, 2018

もっと単純な例だと、以下のような感じです。

import keras

class SimpleMLP(keras.Model):

    def __init__(self, use_bn=False, use_dp=False, num_classes=10):
        super(SimpleMLP, self).__init__(name='mlp')
        self.use_bn = use_bn
        self.use_dp = use_dp
        self.num_classes = num_classes

        self.dense1 = keras.layers.Dense(32, activation='relu')
        self.dense2 = keras.layers.Dense(num_classes, activation='softmax')
        if self.use_dp:
            self.dp = keras.layers.Dropout(0.5)
        if self.use_bn:
            self.bn = keras.layers.BatchNormalization(axis=-1)

    def call(self, inputs):
        x = self.dense1(inputs)
        if self.use_dp:
            x = self.dp(x)
        if self.use_bn:
            x = self.bn(x)
        return self.dense2(x)

model = SimpleMLP()
model.compile(...)
model.fit(...)

注意点として、Modelクラスは__setattr__とかで各レイヤの管理を行っており、直接 self.layer_name = SomeLayer() というコードを書かないとoptimizerから変数が見えなくて全く学習がされていないなどの罠がある。

複数のレイヤをfor文などを利用して上手く定義できたと思ったら上手く動いていないということが起こり得るので注意が必要。

また、現在の実装だとfitの際にモデルが利用するメモリの計算にバグがあるようで、以下の様にサイズを明示してからfitしないとGPUのメモリが枯渇して動作しない。

# batch sizeが1相当のデータを指定
dummy_x = np.zeros((1, 28, 28, 1))
model._set_inputs(dummy_x)

まとめ

とりあえず記法的にはModelを継承してeager_executionを利用しているのが幸せそう。ただし、冒頭に書いた制限やバグもちらほらあるようなので注意は必要。