f:id:Kurene:20210110000321p:plain

Kerasでカスタムレイヤーを含むモデルを保存／読み込むヒントとなる情報があったので試してみました。

まず、以下の日本語のドキュメントの方法のコードでは、Custom layerを含むモデルのsave/loadはできませんでした。

オリジナルのKerasレイヤーを作成する - Keras Documentation

このコードに対して、以下の2点を追加したところsave/load_modelを実行できました。

get_config()の追加
load_model()にキーワード引数custom_objectsを追加

各項目の詳細については、コードや参考Webページをご覧ください。

特にハマったところは、get_config()の書き方でした。

いろいろ検索したところ、以下のWebページに有力そうな情報が書かれていました。

load_model() with custom layers, and custom layers in general · Issue #4871 · keras-team/keras · GitHub

つまり、

カスタムレイヤー＋Layerクラスのコンストラクタ（__init__）のキーワード引数名と値を辞書として返せばよい

とのこと。

以下のソースコードでは、親クラスであるLayerクラスのコンストラクタの引数名・値はbase_configで、カスタムレイヤーでのコンストラクタの引数名・値はconfigに入っており、最終的にそれを結合して返しています。

検証用のソースコード

クリックで展開

from keras import backend as K
from keras.engine.topology import Layer

from keras import Input, Model
import numpy as np


"""
オリジナルのソースコード:
https://keras.io/ja/layers/writing-your-own-keras-layers/
"""

class MyLayer(Layer):
    def __init__(self, output_dim, **kwargs):
        # 外部から与えるパラメタ等を定義
        self.output_dim = output_dim
        super(MyLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        # 学習する重みに関する定義
        self.kernel = self.add_weight(name='kernel',
                                      shape=(input_shape[1], self.output_dim),
                                      initializer='uniform',
                                      trainable=True)
        super(MyLayer, self).build(input_shape)  # Be sure to call this somewhere!

    def call(self, x):
        # Forward　の計算フロー構築
        return K.dot(x, self.kernel)

    def compute_output_shape(self, input_shape):
        # 出力するテンソルの形状を返す
        return (input_shape[0], self.output_dim)
    
    def get_config(self):
        config = {'output_dim': self.output_dim}
        base_config = super(MyLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))
    

n_input  = 8
n_output = 4
np.random.seed(2)
x = np.random.random((5, n_input))
y = np.random.random((5, n_output))

inputs  = Input((n_input, ))
outputs = MyLayer(n_output)(inputs)
model = Model(inputs, outputs)

# トレーニング （今回は無意味な訓練）
model.compile(loss='mean_squared_error',optimizer='sgd')
model.fit(x, y, epochs=10, batch_size=x.shape[0])

y = model.predict(x)
print(f"Before save: model.predict(x) => \n{y}")

from keras.models import load_model
model.save('my_model.h5')
model = load_model('my_model.h5', custom_objects={'MyLayer': MyLayer})

y = model.predict(x)
print(f"After load_model: model.predict(x) => \n{y}")

実行結果

model.save(),model=model_load() の前後でmodel.predict(x) を実施し出力が一致するか検証しました。

なお、トレーニングの内容は適当です。

Epoch 1/10
1/1 [==============================] - 0s 156ms/step - loss: 0.4172
Epoch 2/10
1/1 [==============================] - 0s 0s/step - loss: 0.4125
Epoch 3/10
1/1 [==============================] - 0s 0s/step - loss: 0.4078
Epoch 4/10
1/1 [==============================] - 0s 0s/step - loss: 0.4032
Epoch 5/10
1/1 [==============================] - 0s 0s/step - loss: 0.3987
Epoch 6/10
1/1 [==============================] - 0s 0s/step - loss: 0.3942
Epoch 7/10
1/1 [==============================] - 0s 0s/step - loss: 0.3898
Epoch 8/10
1/1 [==============================] - 0s 0s/step - loss: 0.3855
Epoch 9/10
1/1 [==============================] - 0s 0s/step - loss: 0.3812
Epoch 10/10
1/1 [==============================] - 0s 16ms/step - loss: 0.3770
Before save: model.predict(x) =>
[[ 0.00595104  0.04227125  0.07621967 -0.00854966]
 [-0.02209851  0.05556415  0.06378415 -0.00197074]
 [-0.01329668 -0.0070225   0.05794774 -0.03990902]
 [-0.03630912  0.01880281  0.02773816 -0.01863302]
 [ 0.00922568  0.02720691  0.10599497 -0.02486535]]
After load_model: model.predict(x) =>
[[ 0.00595104  0.04227125  0.07621967 -0.00854966]
 [-0.02209851  0.05556415  0.06378415 -0.00197074]
 [-0.01329668 -0.0070225   0.05794774 -0.03990902]
 [-0.03630912  0.01880281  0.02773816 -0.01863302]
 [ 0.00922568  0.02720691  0.10599497 -0.02486535]]