Keras(TensorFlow)のカスタムレイヤ定義でモデルやロス関数を自由にカスタマイズする

2022.05.30

こんちには。

データアナリティクス事業本部機械学習チームの中村です。

今回は、Kerasで作成できる機械学習モデルの様々な書き方について紹介しつつ、

様々なカスタマイズが可能なカスタムレイヤ定義をご紹介したいと思います。

Kerasについて

もともとKerasは、TensorFlowやその他のライブラリ(Theanoなど)をバックエンドにする、高水準APIライブラリという位置づけでした。

しかし、TensorFlow 2.0以降では、KerasはTensorFlow内のtf.kerasのことを指し、事実上KerasはTensorFlowと同一のものを指すこととなっています。

またモデルをカスタマイズする際であっても、今回紹介するようなカスタムレイヤ定義などを活用すれば、ほとんどのケースにおいてKerasのみで記述可能となっています。

モデル定義方法のパターン

それではKerasでモデルを定義するパターンについて抑えていきます。

Kerasには、大きく分けて3種類のパターンの書き方が存在します。

  • (1) Sequential API
  • (2) Functional API
  • (3) カスタムレイヤ定義

(1)は最も簡易ですが、より複雑なモデルの記述の際には(2),(3)が必要となります。

(2)と(3)ではできることの範囲はほぼ同じなのですが、(2)のままではクラス化ができないため、 より複雑なモデルを構成する場合は、クラス構造で記述できる(3)がより良い選択肢となっています。

ここからはそれぞれの実際の書き方について説明し、ついでに(3)の応用例としてカスタムされたロス関数をどのように実装するかご紹介します。

前準備

モジュールのインポートやサンプルデータの作成を行います。

データは今回ダミーとして乱数を使用します。

import tensorflow as tf
import numpy as np

train_x = np.random.randn(1000, 300)      # 300次元の入力データを1000個生成
train_y = np.random.randint(0, 1+1, 1000) # 0,1のラベルを1000個生成

Sequential API

Sequential APIはもっともシンプルな構築方法で、公式チュートリアルでも一番最初に紹介されています。

Sequential APIはネットワークの分岐や合流がない場合に使用できます。

以下のように記述します。

# 処理される順にSequentialの引数としてレイヤのリストを与える
# この場合、まず中間層の線形層を一つ経由
# その後、分類したいカテゴリ数になるよう処理し、softmaxで確率化
model = tf.keras.Sequential([
    tf.keras.layers.Dense(256, activation='relu'),
    tf.keras.layers.Dense(2, activation='softmax'),
])

# 学習に必要な情報を与える
model.compile(
    optimizer='Adam', 
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# 学習実行
model.fit(train_x, train_y, epochs=5)

学習時のログは以下のようになります。

Epoch 1/5
32/32 [==============================] - 0s 2ms/step - loss: 0.8412 - accuracy: 0.5020
Epoch 2/5
32/32 [==============================] - 0s 2ms/step - loss: 0.4541 - accuracy: 0.8200
Epoch 3/5
32/32 [==============================] - 0s 3ms/step - loss: 0.3209 - accuracy: 0.9240
Epoch 4/5
32/32 [==============================] - 0s 3ms/step - loss: 0.2214 - accuracy: 0.9780
Epoch 5/5
32/32 [==============================] - 0s 2ms/step - loss: 0.1481 - accuracy: 0.9960
<keras.callbacks.History at 0x7f0e35a44dd0>

Functional API

Functional APIはSequential APIと異なり、様々な分岐や結合を表現できます。

以下は途中で分岐させて、結合する例となります。

# 入力はInputクラスで定義する必要がある
x = tf.keras.layers.Input(shape=(300,))

# 分岐させて、それぞれを線形層で処理する
y1 = tf.keras.layers.Dense(64, activation='relu')(x[:,   :100])
y2 = tf.keras.layers.Dense(64, activation='relu')(x[:,100:200])
y3 = tf.keras.layers.Dense(64, activation='relu')(x[:,200:   ])

# 複数の入力を結合する
y = tf.keras.layers.Concatenate()([y1, y2, y3])

# 分類したいカテゴリ数になるよう処理し、softmaxで確率化
y = tf.keras.layers.Dense(2, activation='softmax')(y)

# tf.keras.Modelにネットワークの先頭と終端を与えてモデルを定義する
model = tf.keras.Model(inputs=x, outputs=y)

# 学習に必要な情報を与える(この部分はSequential APIと同じ)
model.compile(
    optimizer='Adam', 
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# 学習実行(この部分はSequential APIと同じ)
model.fit(train_x, train_y, epochs=5)

また入力を最初から分岐させて複数入力することも可能です。

その場合、fitなどでの使用時にも複数の入力を与えてあげる必要があります。

# 複数入力の場合は、Inputクラスを複数定義
x1 = tf.keras.layers.Input(shape=(100,))
x2 = tf.keras.layers.Input(shape=(100,))
x3 = tf.keras.layers.Input(shape=(100,))

y1 = tf.keras.layers.Dense(64, activation='relu')(x1)
y2 = tf.keras.layers.Dense(64, activation='relu')(x2)
y3 = tf.keras.layers.Dense(64, activation='relu')(x3)

y = tf.keras.layers.Concatenate()([y1, y2, y3])

y = tf.keras.layers.Dense(2, activation='softmax')(y)

# 複数入力をinputs引数に与えるように変更
model = tf.keras.Model(inputs=[x1, x2, x3], outputs=y)

model.compile(
    optimizer='Adam', 
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# 学習実行時は、tf.keras.Modelで与えたinputsに合うように与える
model.fit(
    [
        train_x[:,  0:100],
        train_x[:,100:200], 
        train_x[:,200:   ], 
    ],
    train_y, epochs=5
)

カスタムレイヤ定義

カスタムレイヤ定義はFunctional APIと同様に、様々な分岐や結合を表現できます。

クラス構造でオブジェクト指向に記述でき、printなどを使ったデバッグもし易い方法になります。

# カスタムレイヤクラスを定義
class CustomModel(tf.keras.Model):

    # パラメータを持つレイヤは、__init__でクラスメンバとして定義する
    def __init__(self):
        super(CustomModel, self).__init__()

        self.dense1 = tf.keras.layers.Dense(64, activation='relu')
        self.dense2 = tf.keras.layers.Dense(64, activation='relu')
        self.dense3 = tf.keras.layers.Dense(64, activation='relu')
        self.dense_final = tf.keras.layers.Dense(2, activation='softmax')

    # __call__では、fit時などに実際にモデルで処理する場合の操作を定義する
    def call(self, input_tensor, training=False):

        y1 = self.dense1(input_tensor[:,   :100], training=training)
        y2 = self.dense2(input_tensor[:,100:200], training=training)
        y3 = self.dense3(input_tensor[:,200:   ], training=training)

        y = tf.keras.layers.Concatenate()([y1, y2, y3])

        return self.dense_final(y)

# カスタムレイヤクラスを実体化
model = CustomModel()

# これ以降は今までと同様
model.compile(
    optimizer='Adam', 
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

model.fit(train_x, train_y, epochs=5)

上記の例は、Functional APIの一つ目の例と同じ構成となります。

カスタムレイヤを定義する際に必須なのは、__init____call__の2つのメソッドになります。

__call__部分にprint処理などを挿入することで、途中の値などをデバッグすることができます。

また__call__には、trainingという引数がありますが、これは学習時と推論時で処理を分岐する必要がある際に使用します。

lossおよびmetricのカスタマイズ

カスタムレイヤでは、__call__の中でself.add_lossを使えば、compile時に指定しているロスに加算することができます。

またself.add_metricを使えば、学習時のログに表示されるmetricを追加できます。

class CustomModel(tf.keras.Model):

    def __init__(self):
        super(CustomModel, self).__init__()

        self.dense1 = tf.keras.layers.Dense(64, activation='relu')
        self.dense2 = tf.keras.layers.Dense(64, activation='relu')
        self.dense3 = tf.keras.layers.Dense(64, activation='relu')
        self.dense_final = tf.keras.layers.Dense(2, activation='softmax')

    def call(self, input_tensor, training=False):

        y1 = self.dense1(input_tensor[:,   :100], training=training)
        y2 = self.dense2(input_tensor[:,100:200], training=training)
        y3 = self.dense3(input_tensor[:,200:   ], training=training)

        y = tf.keras.layers.Concatenate()([y1, y2, y3])

        # lossに加算(今回は適当なダミー定数)
        self.add_loss(100)

        # metricに追加(今回は適当なダミー定数)
        self.add_metric(100, name="test")

        return self.dense_final(y)

model = CustomModel()

model.compile(
    optimizer='Adam', 
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

model.fit(train_x, train_y, epochs=5)

この場合、ログは以下のようになります。

Epoch 1/5
32/32 [==============================] - 1s 2ms/step - loss: 100.9088 - accuracy: 0.4740 - test: 100.0000
Epoch 2/5
32/32 [==============================] - 0s 2ms/step - loss: 100.6924 - accuracy: 0.6040 - test: 100.0000
Epoch 3/5
32/32 [==============================] - 0s 2ms/step - loss: 100.5878 - accuracy: 0.6850 - test: 100.0000
Epoch 4/5
32/32 [==============================] - 0s 2ms/step - loss: 100.5104 - accuracy: 0.7660 - test: 100.0000
Epoch 5/5
32/32 [==============================] - 0s 2ms/step - loss: 100.4454 - accuracy: 0.8340 - test: 100.0000
<keras.callbacks.History at 0x7f0e30784790>

ロスに100が加算されており、testというmetricが追加されていることがわかります。

add_lossが指定されている場合、compileのloss引数は場合によっては削除することも可能です。

(今回のように定数をadd_lossしている場合はパラメータが更新できないためエラーとなります)

つまり、カスタムレイヤの方のadd_lossで最終的なロスを記述すれば、完全にカスタムな損失関数を使えます。

なおadd_lossadd_metricは、Functional APIでも使用可能ですが、今回は省略します。

まとめ

いかがでしたでしょうか。カスタムレイヤのより詳しい説明はTensorFlow公式ページにもありますので、そちらもご覧いただければと思います。

この記事が実際に自作モデルを構築する際の一助となりましたら幸いです。