KerasのCustom Layerを保存・ロードする方法｜get_configと登録を解説

KerasのCustom Layerを.keras形式で保存し、同じモデルへ復元するイメージ

KerasのCustom Layerを.keras形式で保存し、同じモデルへ復元するイメージ

Keras 3でCustom Layer（独自に実装したレイヤー）を含むモデルを保存・ロードするなら、基本は次の4点です。

  1. Custom Layerの__init__()で受け取る設定をget_config()へ含める
  2. @keras.saving.register_keras_serializableでクラスを登録する
  3. モデル全体を.keras形式で保存する
  4. ロード後にレイヤーの型、設定、推論出力を確認する

登録デコレータを使えば、通常はload_model()へcustom_objectsを毎回渡す必要がありません。ただし、.kerasファイルにはCustom LayerのPythonコード自体は入りません。ロード側でもクラス定義をimportし、デコレータの登録処理を実行しておく必要があります。

この記事では、Keras 3のkeras.saving APIへ表記を統一し、Custom Layerの定義から保存・復元・確認までを一つの例で説明します。

.kerasに保存されるものと、保存されないもの

Keras公式の保存ガイドによると、.keras形式には主に次の情報が保存されます。

  • モデルの構成
  • 学習済みの重み
  • compile()済みなら損失関数やメトリクスなどのコンパイル情報
  • オプティマイザとその状態

一方、Custom Layerを定義したPythonコードそのものは保存されません。Kerasは保存された構成からクラスを探し、get_config()の内容を使ってインスタンスを作り直した後、重みを戻します。

要素 .kerasに保存されるか ロード側で必要なもの
レイヤーの構成 保存される get_config()から再構築できること
学習済み重み 保存される 対応するレイヤー構造
Custom LayerのPythonコード 保存されない クラス定義をimportして利用可能にする
Custom Objectの名前解決 登録名が構成へ記録される 登録済みクラス、またはcustom_objects

KerasのCustom Layerでconfigとweightsを保存し、ロード時にクラス登録を使って復元する流れ

モデル全体を後からそのまま使いたい場合は.kerasが向きます。save_weights()は重みだけを保存するため、ロード時に同じモデル構造をコードで作り直せる場合に使います。

Custom Layerを保存・ロードする最小例

次の例では、入力の各要素に学習可能なscaleとbiasを適用するAffineLayerを作ります。initial_scaleはレイヤーを再構築するための設定値、scaleとbiasは保存対象の重みです。

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

import keras
import numpy as np

LOGGER = logging.getLogger(__name__)
MODEL_PATH = Path("affine_model.keras")


@keras.saving.register_keras_serializable(package="ExampleLayers")
class AffineLayer(keras.layers.Layer):
    """入力へ要素ごとのアフィン変換を適用するレイヤー。

    Args:
        initial_scale: scaleの初期値。
        **kwargs: `name`や`dtype`など、基底Layerへ渡す設定。

    Returns:
        AffineLayerのインスタンス。

    Raises:
        TypeError: `initial_scale`をfloatへ変換できない場合。

    Example:
        >>> layer = AffineLayer(initial_scale=2.0, name="affine")
    """

    def __init__(self, initial_scale: float = 1.0, **kwargs: Any) -> None:
        """レイヤーを初期化する。

        Args:
            initial_scale: scaleの初期値。
            **kwargs: 基底Layerへ渡すキーワード引数。

        Returns:
            None。

        Raises:
            TypeError: `initial_scale`をfloatへ変換できない場合。

        Example:
            >>> AffineLayer(initial_scale=0.5)
        """
        super().__init__(**kwargs)
        self.initial_scale = float(initial_scale)
        self.scale: Any = None
        self.bias: Any = None

    def build(self, input_shape: tuple[int | None, ...]) -> None:
        """入力の最終次元に対応する重みを作る。

        Args:
            input_shape: バッチ次元を含む入力形状。

        Returns:
            None。

        Raises:
            ValueError: 入力の最終次元が未確定の場合。

        Example:
            >>> layer = AffineLayer()
            >>> layer.build((None, 3))
        """
        feature_count = input_shape[-1]
        if feature_count is None:
            raise ValueError("入力の最終次元を確定してください。")

        self.scale = self.add_weight(
            name="scale",
            shape=(feature_count,),
            initializer=keras.initializers.Constant(self.initial_scale),
            trainable=True,
        )
        self.bias = self.add_weight(
            name="bias",
            shape=(feature_count,),
            initializer="zeros",
            trainable=True,
        )
        super().build(input_shape)

    def call(self, inputs: Any) -> Any:
        """入力へscaleとbiasを適用する。

        Args:
            inputs: 最終次元がbuild時の次元と一致するテンソル。

        Returns:
            `inputs * scale + bias`の結果。

        Raises:
            ValueError: 入力形状がレイヤーの重みと一致しない場合。

        Example:
            >>> layer = AffineLayer(initial_scale=2.0)
            >>> _ = layer(keras.ops.ones((1, 3)))
        """
        return inputs * self.scale + self.bias

    def get_config(self) -> dict[str, Any]:
        """レイヤーを再構築する設定を返す。

        Args:
            なし。

        Returns:
            基底Layerの設定と`initial_scale`を含む辞書。

        Raises:
            RuntimeError: 基底Layerが設定を返せない場合。

        Example:
            >>> config = AffineLayer(2.0).get_config()
            >>> config["initial_scale"]
            2.0
        """
        config = super().get_config()
        config.update({"initial_scale": self.initial_scale})
        return config


def build_model() -> keras.Model:
    """AffineLayerを含むFunctionalモデルを作る。

    Args:
        なし。

    Returns:
        3要素の入力を同じ形状で変換するKerasモデル。

    Raises:
        ValueError: レイヤーの入出力形状を構築できない場合。

    Example:
        >>> model = build_model()
        >>> model.output_shape
        (None, 3)
    """
    inputs = keras.Input(shape=(3,), name="features")
    outputs = AffineLayer(initial_scale=2.0, name="affine")(inputs)
    return keras.Model(inputs=inputs, outputs=outputs, name="affine_model")


def main() -> None:
    """モデルを保存・ロードし、復元結果を検証する。

    Args:
        なし。

    Returns:
        None。

    Raises:
        AssertionError: ロード前後の型、設定、推論値が一致しない場合。
        OSError: モデルファイルを保存または読み込めない場合。

    Example:
        >>> main()
    """
    model = build_model()
    sample = np.array([[1.0, 2.0, 3.0]], dtype=np.float32)
    output_before = keras.ops.convert_to_numpy(model(sample, training=False))

    model.save(MODEL_PATH)
    LOGGER.info("モデルを保存しました: %s", MODEL_PATH)

    restored = keras.saving.load_model(MODEL_PATH, compile=False)
    output_after = keras.ops.convert_to_numpy(restored(sample, training=False))

    restored_layer = restored.get_layer("affine")
    assert isinstance(restored_layer, AffineLayer)
    assert restored_layer.initial_scale == 2.0
    np.testing.assert_allclose(output_before, output_after, rtol=1e-6, atol=1e-6)
    LOGGER.info("Custom Layerの設定と推論値が一致しました")


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    main()

実行にはKeras 3と、TensorFlow・JAX・PyTorchのいずれかのバックエンドが必要です。たとえばTensorFlowバックエンドを使う環境では、必要なパッケージをインストールしたうえで実行します。プロジェクトでは依存バージョンを固定し、保存側とロード側で同じ組み合わせを再現できるようにしてください。

成功時はaffine_model.kerasが作られ、ロード後のレイヤーがAffineLayerであること、initial_scaleが2.0であること、推論値が許容誤差内で一致することを確認できます。

推論結果へ.numpy()を直接呼ばず、keras.ops.convert_to_numpy()を使っているのは、TensorFlow・JAX・PyTorchで返されるテンソル型の差を呼び出し側へ持ち込まないためです。Custom Layerの計算もkeras.opsで表現すると、特定バックエンドの演算へ依存しにくくなります。

get_config()はコンストラクタ設定を保存する

get_config()が保存するのは、レイヤーを作り直すための設定です。学習で更新されるscaleやbiasの値を手動で辞書へ入れる必要はありません。add_weight()で作成したtrainable/non-trainable weightsは、Kerasの通常の重み保存処理で追跡されます。

def get_config(self) -> dict[str, Any]:
    """レイヤーの再構築に必要な設定を返す。

    Args:
        なし。

    Returns:
        基底Layerの設定と独自設定を含む辞書。

    Raises:
        RuntimeError: 基底Layerが設定を返せない場合。

    Example:
        >>> AffineLayer(2.0).get_config()["initial_scale"]
        2.0
    """
    config = super().get_config()
    config.update({"initial_scale": self.initial_scale})
    return config

super().get_config()を使うのは、name、trainable、dtypeなど、基底Layerが管理する設定も残すためです。対応して、__init__()はkwargsを受け取り、super().__init__(kwargs)へ渡します。

initial_scaleはfloatなので、既定のfrom_config()がAffineLayer(**config)を実行すれば復元できます。このようにコンストラクタ引数がJSONで表現できる単純値だけなら、from_config()を上書きする必要はありません。

登録デコレータはクラスの検索名を決める

@keras.saving.register_keras_serializable(package="ExampleLayers")
class AffineLayer(keras.layers.Layer):
    ...

Serialization utilitiesによると、登録名はpackage>name形式です。この例では、名前を省略しているためExampleLayers>AffineLayerとして登録されます。

保存ファイルに登録名があっても、クラスの実装コードが埋め込まれるわけではありません。クラスをmy_project/layers.pyへ分けた場合は、ロード前にそのモジュールをimportします。

from my_project.layers import AffineLayer  # デコレータを実行して登録するために必要
import keras

model = keras.saving.load_model("affine_model.keras", compile=False)

AffineLayerをこのファイルで直接使っていなくても、importにはレジストリへ登録する役割があります。パッケージ名や登録名を後から変更すると、以前のモデルが古い名前を参照できなくなるため、公開済みモデルを維持するプロジェクトでは登録名を互換性のあるAPIとして扱います。

Custom Objectを解決する3つの方法

Kerasは登録以外にも、ロード時の辞書と一時的なscopeを提供しています。

方法 向く場面 長所 トレードオフ
register_keras_serializable 自分で保守するCustom Layer 呼び出し側が簡潔で、保存形式上の名前も安定させやすい ロード前に定義モジュールのimportが必要。登録名の変更は互換性へ影響する
custom_objects 既存クラスを変更できない、ロード箇所が少ない 依存関係がロード行に明示される すべてのロード箇所で辞書を渡す必要がある
custom_object_scope 複数回のデシリアライズを限定範囲で行う scopeの内側だけ一時的に解決できる scopeの外では使えず、処理範囲が広いと依存が見えにくい

custom_objectsを渡す

デコレータを付けられないクラスでは、保存時に使った名前とクラスの対応を渡します。

restored = keras.saving.load_model(
    "affine_model.keras",
    custom_objects={"AffineLayer": AffineLayer},
    compile=False,
)

登録名付きで保存されたファイルを扱う場合は、エラーに表示されるregistered_nameと対応するキーも確認します。保存時のKerasバージョンや実際のconfigによって検索される名前が異なるため、推測でキーを増やすより、登録方式を統一するか、保存ファイルのエラー表示を確認します。

custom_object_scope()で範囲を限定する

custom_objects = {"AffineLayer": AffineLayer}

with keras.saving.custom_object_scope(custom_objects):
    restored = keras.saving.load_model("affine_model.keras", compile=False)

scopeを抜けると、この辞書による一時登録は使われません。テストや移行コードなど、Custom Objectの解決範囲を明示したい場合に向きます。

レイヤーや関数を引数に持つ場合はfrom_config()も実装する

Custom Layerのコンストラクタが別のレイヤーを受け取る場合、そのオブジェクトをそのままJSONへ保存することはできません。Keras公式ガイドは、serialize_keras_object()とdeserialize_keras_object()を使う方法を示しています。

from typing import Any

import keras


@keras.saving.register_keras_serializable(package="ExampleLayers")
class WrappedLayer(keras.layers.Layer):
    """別のKerasレイヤーを内包して呼び出すレイヤー。

    Args:
        inner_layer: 内部で実行するKerasレイヤー。
        **kwargs: 基底Layerへ渡す設定。

    Returns:
        WrappedLayerのインスタンス。

    Raises:
        TypeError: `inner_layer`がKeras Layerではない場合。

    Example:
        >>> WrappedLayer(keras.layers.Dense(4))
    """

    def __init__(self, inner_layer: keras.layers.Layer, **kwargs: Any) -> None:
        """内部レイヤーを保持する。

        Args:
            inner_layer: 内部で実行するKerasレイヤー。
            **kwargs: 基底Layerへ渡す設定。

        Returns:
            None。

        Raises:
            TypeError: `inner_layer`がKeras Layerではない場合。

        Example:
            >>> WrappedLayer(keras.layers.Dense(4))
        """
        if not isinstance(inner_layer, keras.layers.Layer):
            raise TypeError("inner_layerにはKeras Layerを指定してください。")
        super().__init__(**kwargs)
        self.inner_layer = inner_layer

    def call(self, inputs: Any) -> Any:
        """内部レイヤーへ入力を渡す。

        Args:
            inputs: 内部レイヤーが受け取れるテンソル。

        Returns:
            内部レイヤーの出力。

        Raises:
            ValueError: 内部レイヤーが入力を処理できない場合。

        Example:
            >>> layer = WrappedLayer(keras.layers.Dense(4))
            >>> _ = layer(keras.ops.ones((1, 3)))
        """
        return self.inner_layer(inputs)

    def get_config(self) -> dict[str, Any]:
        """内部レイヤーを含む設定をシリアライズする。

        Args:
            なし。

        Returns:
            内部レイヤーのシリアライズ結果を含む辞書。

        Raises:
            TypeError: 内部レイヤーをシリアライズできない場合。

        Example:
            >>> layer = WrappedLayer(keras.layers.Dense(4))
            >>> "inner_layer" in layer.get_config()
            True
        """
        config = super().get_config()
        config.update(
            {"inner_layer": keras.saving.serialize_keras_object(self.inner_layer)}
        )
        return config

    @classmethod
    def from_config(cls, config: dict[str, Any]) -> "WrappedLayer":
        """保存された設定からレイヤーを復元する。

        Args:
            config: `get_config()`が返した設定。

        Returns:
            復元したWrappedLayer。

        Raises:
            TypeError: 内部レイヤーの設定を復元できない場合。

        Example:
            >>> original = WrappedLayer(keras.layers.Dense(4))
            >>> restored = WrappedLayer.from_config(original.get_config())
        """
        copied_config = dict(config)
        inner_config = copied_config.pop("inner_layer")
        inner_layer = keras.saving.deserialize_keras_object(inner_config)
        return cls(inner_layer=inner_layer, **copied_config)

dict(config)でコピーしているのは、呼び出し元が持つ設定辞書をpop()で変更しないためです。内部オブジェクトもCustom Objectなら、そのクラスも登録するか、ロード時に解決できるようにします。

get_config()だけでは足りないケース

通常のCustom Layerでは、コンストラクタ設定をget_config()へ入れ、重みをadd_weight()で作れば足ります。保存対象がその範囲を超える場合は、Customizing Saving and Serializationで説明されている拡張フックを使います。

保存したいもの 主なAPI 注意点
add_weight()で追跡されない独自変数 save_own_variables() / load_own_variables() 上書きした場合は通常の重みも自分で保存・復元する
語彙、ラベル一覧などのファイル save_assets() / load_assets() ファイル形式と文字コードの互換性を管理する
通常と異なるbuild引数や状態 get_build_config() / build_from_config() ロード時に同じ重み形状を再構築できるようにする
独自compile情報 get_compile_config() / compile_from_config() 推論だけならload_model(..., compile=False)で不要な復元を避けられる

これらはget_config()の代替ではありません。モデル構成、重み、アセットなど、保存したい情報の種類ごとに役割を分けます。

よくあるエラーと直し方

Could not locate class 'AffineLayer'

Custom Layerがロード時のレジストリにありません。

  • クラスへregister_keras_serializableを付けたか
  • ロード前にクラス定義のモジュールをimportしたか
  • 保存後にpackageやnameを変更していないか
  • 登録しない方式ならcustom_objectsを渡したか

同名クラスを別ファイルへ複製して回避するより、保存側とロード側が参照する定義を一つにします。

__init__()に予期しないnameやtrainableが渡される

super().get_config()には基底Layerの設定も含まれます。Custom Layerの__init__()でkwargsを受け取り、super().__init__(kwargs)へ渡してください。

def __init__(self, initial_scale: float = 1.0, **kwargs: Any) -> None:
    """基底Layerの設定を受け取って初期化する。

    Args:
        initial_scale: scaleの初期値。
        **kwargs: `name`、`dtype`、`trainable`などの設定。

    Returns:
        None。

    Raises:
        TypeError: 基底Layerが引数を受け取れない場合。

    Example:
        >>> AffineLayer(initial_scale=1.0, name="affine")
    """
    super().__init__(**kwargs)
    self.initial_scale = float(initial_scale)

ネストしたレイヤーを正しく復元できない

get_config()へレイヤーオブジェクトをそのまま入れず、serialize_keras_object()で辞書へ変換します。ロード側ではfrom_config()でdeserialize_keras_object()を呼び、実体へ戻してからコンストラクタへ渡します。

ロードは成功するが推論値が違う

保存前後で、次の条件をそろえます。

  • 同じ入力を使う
  • 推論確認ではtraining=Falseにする
  • np.testing.assert_allclose()で浮動小数点の許容誤差を指定する
  • 保存前に重みが作成済みか確認する
  • 独自変数をadd_weight()以外で持つ場合、専用の保存フックを実装する
  • 保存側とロード側のKeras、バックエンド、独自コードのバージョンを記録する

変数名はロード後に変わる可能性があるため、文字列の変数名に依存せず、get_layer()やレイヤー属性から取得します。Whole model saving & loadingも、特定の変数へはレイヤー属性でアクセスすることを勧めています。

信頼できない.kerasファイルをロードしない

Keras 3のload_model()はsafe_mode=Trueが既定です。これは危険なlambdaのデシリアライズを禁止しますが、Pythonプロセス全体を隔離するsandboxではありません。

restored = keras.saving.load_model(
    "affine_model.keras",
    compile=False,
    safe_mode=True,
)

safe_mode=Falseへ変えると、任意コード実行につながる可能性があると公式APIは警告しています。safe modeを有効にしていても、出所不明のモデル、改ざんの可能性があるモデル、信頼できないCustom Object実装を通常の開発環境でロードしないでください。

まとめ

Keras 3でCustom Layerを含むモデルを保存・ロードする基本手順は、get_config()で再構築用の設定を返し、register_keras_serializableでクラス名を解決できるようにして、.keras形式で保存することです。

単純値だけなら既定のfrom_config()で足ります。内部レイヤーなどを引数に持つ場合は、serialize_keras_object()とdeserialize_keras_object()を組み合わせます。ロード後は「エラーが出なかった」だけで終えず、レイヤー型、設定値、推論出力を確認すると、構成と重みの両方を復元できたか判断できます。

出典

スポンサーリンク

コメント

タイトルとURLをコピーしました