みなさんこんにちは。たかぱい（@takapy0210）です。

本日はTensorFlow×Transformers周りでエラーに遭遇した内容とそのWAです。

環境

実行環境は以下の通りです

python 3.7.10
transformers 4.12.5
tensorflow 2.3.0

実装内容

一部抜粋ですが、TransformersのTFBertModel*1に、独自のレイヤーをいくつか追加した2値分類モデルの学習を行いました。

import tensorflow as tf
from tensorflow import keras
from transformers import TFBertModel


def build_model(model_name, num_classes, max_length):
    
    # BERTへ入力する情報
    input_shape = (max_length, )
    input_ids = tf.keras.layers.Input(input_shape, dtype=tf.int32)
    attention_mask = tf.keras.layers.Input(input_shape, dtype=tf.int32)
    token_type_ids = tf.keras.layers.Input(input_shape, dtype=tf.int32)
    inputs = [input_ids, attention_mask, token_type_ids]
    
    # BERTモデル
    bert_model = TFBertModel.from_pretrained(model_name)
    # 2種類の出力がある。今回使用するのはTFBertForSequenceClassificationにならってpooler_outputを使用
    base_output = bert_model(inputs)
    sequence_output, pooler_output = base_output[0], base_output[1]
    
    output = tf.keras.layers.Dense(num_classes, activation="sigmoid")(pooler_output)
    model = tf.keras.Model(inputs=inputs, outputs=output)
    optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0)
    model.compile(optimizer=optimizer, loss="binary_crossentropy", metrics=["acc"])
    return model


# データの前処理
# （省略）
...

# 学習
model.fit(
    X_train,
    y_train,
    validation_data=(X_valid, y_valid),
    batch_size=batch_size,
    epochs=epochs,
    callbacks=callbacks
)

model.save('bert_model.h5')

エラー内容

model.save('bert_model.h5') の部分で以下のようなエラーが発生します。

（省略）
/usr/local/lib/python3.7/site-packages/transformers/modeling_tf_utils.py in input_processing(func, config, input_ids, **kwargs)
    408                     output[tensor_name] = input
    409                 else:
--> 410                     output[parameter_names[i]] = input
    411             elif isinstance(input, allowed_types) or input is None:
    412                 output[parameter_names[i]] = input

IndexError: list index out of range

エラーの原因

コードを見てみると、inputsとして配列が想定されていないからかな〜と思います（多分）

github.com

ワークアラウンド

BERTモデルを読み込む部分の処理を以下のように修正すればOKです。

# エラーになるコード
base_output = bert_model(inputs)

# WA
base_output = bert_model.bert(inputs)

なんでこれで解消できるのか？

実装されているコード*2をみると、bert layerというのは TFBertMainLayer*3 というクラスです。
これを直接呼ぶことで、Save時も TFBertMainLayer.call が呼び出され、LayerCall.__call__(args, **kwargs) が呼び出されることを回避することができるので、保存処理がうまくいくのだと思います。

モデルの保存方法

保存するときはcallback関数に設定しても、model.save() メソッドを用いても、どちらでも問題なく保存できました。

# callback関数を使う場合
checkpoint = tf.keras.callbacks.ModelCheckpoint(
    'bert_model.h5',
    monitor='val_loss',
    verbose=2,
    save_best_only=True,
    mode='auto'
)
result = model.fit(
    X_train,
    y_train,
    validation_data=(X_valid, y_valid),
    batch_size=batch_size,
    epochs=epochs,
    callbacks=[checkpoint]
)

# saveメソッドで保存する場合
model.save('bert_model.h5')

保存したモデルを読み込む際は普通に load_model() を使えば読み込めます。

from tensorflow.keras.models import load_model

model = load_model('bert_model.h5')
pred = model.predict(test, verbose=1)

参考

github.com

*1:https://huggingface.co/docs/transformers/model_doc/bert#transformers.TFBertModel

*2:https://github.com/huggingface/transformers/blob/32e94cff64ea87cb2df2699bca960962fe676b62/src/transformers/modeling_tf_bert.py#L699

*3:https://github.com/huggingface/transformers/blob/32e94cff64ea87cb2df2699bca960962fe676b62/src/transformers/modeling_tf_bert.py#L494

ギークなエンジニアを目指す男

機械学習系の知識を蓄えようとするブログ

TensorFlow × HuggingFace Transformers（TFBertModel）を用いたモデルの保存時に発生するエラーの回避方法

環境