ViTモデルロード時のエラー（ValueError: Unknown layer: ClassToken）について

画像分類モデルViT（Vision Transformer）で学習したモデルをロードしようとしたところエラーが出ました。対処方法についてまとめます。ちなみにViTについてはこちらでまとめてます。機械学習フレームワークはKerasで構築しています。

1. エラー内容
2. 対処方法
3. 試行錯誤（長いので忙しい人は無視してください）

1. エラー内容

以下を実行。

# 学習したモデルを以下で保存
model.save("model.h5", include_optimizer=False)

# 保存したモデルを以下で普通にロード
model = tf.keras.models.load_model('model.h5',compile=False)

すると、以下のエラーが出る。

ValueError: Unknown layer: ClassToken

2. 対処方法

学習時に使用したモデルを定義し直して、重みをロードすればOK。

# 学習時に使用したモデルを再定義
vit_model = vit.vit_b16(
    image_size = image_size_vitb16,
    activation = 'sigmoid',
    pretrained = True,
    include_top = False,
    pretrained_top = False)
model = tf.keras.Sequential([
    vit_model,
    tf.keras.layers.Flatten(),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Dense(11, activation = tfa.activations.gelu),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Dense(num_classes, 'softmax')
])
# ロードするモデルを指定。
model .load_weights('model.h5')

3. 試行錯誤（長いので忙しい人は無視してください）

試行錯誤①：単純にvit-kerasをインポート

エラー文言をそのまま読むと、「Class TokenっていうLayerが定義されていないよ！」って意味になります。たぶん、Class TokenというのはViT独自のレイヤーだと思われます。vit-kerasのソースコードをGitHubで確認したところ、Class TokenというLayerが確かに定義されていました。なので、vit-kerasをとりあえずインポートすればClass Tokenも使えるようになるかなと思い、以下のコードを実行しました。

import vit_keras

しかし、解決せず。
ちなみにvit-kerasのソースコード（GitHub）については以下を参考にしてください。
vit-keras/layers.py at master · faustomorales/vit-keras · GitHub

試行錯誤②：モデルロード時にcustom_objectとしてClass Tokenを定義

次にロードするときにClass Tokenレイヤーを定義するようにソース修正。

from vit_keras import layers
model = tf.keras.models.load_model('model.h5',custom_objects={'ClassToken': layers.ClassToken},compile=False)

結果は以下の通り。

ValueError: Unknown layer: AddPositionEmbs

エラー内容が変わった！vit-kerasのソースを見直してみるとAddPositionEmbsというLayerも発見。
layers.pyに定義されているレイヤーを一通り定義するように修正。

from vit_keras import layers
model = tf.keras.models.load_model('model.h5',
                                   custom_objects={'ClassToken': layers.ClassToken,
                                                   'AddPositionEmbs': layers.AddPositionEmbs,
                                                   'MultiHeadSelfAttention': layers.MultiHeadSelfAttention,
                                                   'TransformerBlock': layers.TransformerBlock},
                                   compile=False)

結果は以下の通り。

TypeError: __init__() missing 3 required keyword-only arguments: 'num_heads', 'mlp_dim', and 'dropout'

TransformerBlockレイヤーを確認したところコンストラクタに、'num_heads', 'mlp_dim', 'dropout'が引数として定義されている。モデルロード時にこれらの引数が渡せていない模様。色々調べたのですが、引数の渡し方が分からず、custom_obejectを使用した方法は断念。