Colabのkerasとtransformers最新版で出る変なエラーへの対処メモ

情報が少ないのでメモを残しておく。

1. 問題の再現

最近いよいよスパコンで計算することが多くなってきてColabはあんまり触ってなかったんだけど、ちょっと個人的なプロジェクトのために触ってみたら変なエラーが出た。
2024年4月26日現在、ColabでKeras + transformersでBERTあたりを組み込んだモデルを構築する場合、デフォルトの環境は
- tensorflow: 2.15.0
- transformers: 4.40.0
になっている。

import dumpy as np
import pandas as pd

import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Input

from transformers import BertTokenizer, TFAutoModel

# Allocate GPU
device_name = tf.test.gpu_device_name()
if device_name != '/device:GPU:0':
  raise SystemError('GPU device not found')
print('GPU: {}'.format(device_name))

# Load BERT Component
model_id = 'bert-base-uncased'
bert_token = BertTokenizer.from_pretrained(model_id)
bert_model = TFAutoModel.from_pretrained(model_id)

こんな感じで準備をしておいて、以下のように本当に簡単な、BERTからのembeddingsを取得するようなコードを書くと、Transformersのレイヤーからkerasへのテンソルの受け渡しにあたって結構奇怪なエラーが出る場合がある。

maxlen = 256
shape = (maxlen, )
ids = Input(shape=shape, dtype = tf.int32, name = "input_id")
att = Input(shape=shape, dtype = tf.int32, name = "input_att_mask")
tok = Input(shape=shape, dtype = tf.int32, name = "input_token")
x1out = bert_model(input_ids=ids, attention_mask=att, token_type_ids=tok)
x1out = x1out.pooler_output
x1in = [ids, att, tok]

みてわかる通り、6行目っていうのはtransformersのTFAutoModelで作ったBERTのコンポーネントから吐き出したembeddingsをx1outに割り当てているだけの作業なんだけど、そこで以下のようなエラーが出る。

TypeError: Exception encountered when calling layer 'embeddings' (type TFBertEmbeddings).

Could not build a TypeSpec for name: "tf.debugging.assert_less_1/assert_less/Assert/Assert"
op: "Assert"
input: "tf.debugging.assert_less_1/assert_less/All"
input: "tf.debugging.assert_less_1/assert_less/Assert/Assert/data_0"
input: "tf.debugging.assert_less_1/assert_less/Assert/Assert/data_1"
input: "tf.debugging.assert_less_1/assert_less/Assert/Assert/data_2"
input: "Placeholder"
input: "tf.debugging.assert_less_1/assert_less/Assert/Assert/data_4"
input: "tf.debugging.assert_less_1/assert_less/y"
attr {
  key: "T"
  value {
    list {
      type: DT_STRING
      type: DT_STRING
      type: DT_STRING
      type: DT_INT32
      type: DT_STRING
      type: DT_INT32
    }
  }
}
attr {
  key: "summarize"
  value {
    i: 3
  }
}
 of unsupported type <class 'tensorflow.python.framework.ops.Operation'>.

Call arguments received by layer 'embeddings' (type TFBertEmbeddings):
  • input_ids=<KerasTensor: shape=(None, 256) dtype=int32 (created by layer 'input_id')>
  • position_ids=None
  • token_type_ids=<KerasTensor: shape=(None, 256) dtype=int32 (created by layer 'input_token')>
  • inputs_embeds=None
  • past_key_values_length=0
  • training=False

まあ非常に簡単にいえば、入力として受け入れるレイヤーの型(type)があってないってだけなんだけど、まだx1outの出力自体はいじってないわけだし、つまりはBERTへの入力側としてのinput_ids, attention_mask, token_type_idsの型が変だっていう感じっぽい。

個人的には、tensorflowのモデルはcompileするまで構造上のエラーに気づかない（もっというとcompile()後にsummary()するまで気づかない）ことが結構ある印象だけど、まあ今回はそもそもレイヤー同士が繋がらないんだから仕方ないかな。

2. 解決方法

transformersのインストールにあたり明示的にバージョン4.31.0を指定すること。それだけ。pipで入れるときに指定してあげること。

!pip install -q transformers==4.31.0

これでOK。
ただし、すでにインスタンスを立ち上げて一旦transformersの最新版をimportしてしまっている場合には、一旦インスタンスの再起動が必要です。再起動せずにpip installしてimport transformersしても同じエラーが出ます。