TensorFlow 2.15.0中TFBertEmbeddings层TypeError问题求助

问题场景 #

两周前模型可正常运行，现在调用BERT模型时触发以下TypeError，错误指向embeddings层，提示无法为Assert操作构建TypeSpec：

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-23-a3e5a45f06c9> in <cell line: 14>()
     12 
     13 # Encode input using BERT model
---> 14 bert_output = bert_model(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
     15 
     16 # Get pooled output and pass through dropout layer

8 frames
/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py in error_handler(*args, **kwargs)
     68             # To get the full stack trace, call:
     69             # `tf.debugging.disable_traceback_filtering()`
---> 70             raise e.with_traceback(filtered_tb) from None
     71         finally:
     72             del filtered_tb

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

Could not build a TypeSpec for name: "tf.debugging.assert_less/assert_less/Assert/Assert"
op: "Assert"
input: "tf.debugging.assert_less/assert_less/All"
input: "tf.debugging.assert_less/assert_less/Assert/Assert/data_0"
input: "tf.debugging.assert_less/assert_less/Assert/Assert/data_1"
input: "tf.debugging.assert_less/assert_less/Assert/Assert/data_2"
input: "Placeholder"
input: "tf.debugging.assert_less/assert_less/Assert/Assert/data_4"
input: "tf.debugging.assert_less/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, 50) dtype=int32 (created by layer 'input_ids')>
  • position_ids=None
  • token_type_ids=<KerasTensor: shape=(None, 50) dtype=int32 (created by layer 'token_type_ids')>
  • inputs_embeds=None
  • past_key_values_length=0
  • training=False

降级TensorFlow旧版本可解决该错误，但会丢失GPU支持且出现Graph Execution Error。

解决方案 #

1. 升级Transformers库到兼容版本 #

TensorFlow 2.15.0对图执行的断言操作处理逻辑有更新，旧版本transformers库未适配导致报错。升级到transformers>=4.35.0即可修复兼容性问题：

pip install --upgrade transformers

若升级后仍有问题，可指定稳定兼容版本：

pip install transformers==4.36.2

2. 临时切换到Eager执行模式绕过图断言问题 #

如果暂时无法升级transformers，可在代码开头添加以下代码，强制TensorFlow以Eager模式执行，绕过图执行中的断言类型错误：

import tensorflow as tf
tf.config.run_functions_eagerly(True)

注意：该方式会降低模型运行性能，仅作为临时验证方案，优先推荐升级transformers。

3. 验证输入张量的合法性 #

确保输入的input_ids、attention_mask、token_type_ids完全符合BERT要求：

• 张量 dtype 为int32
• 形状为(batch_size, seq_len)，无异常维度
• 可通过以下代码检查：

print("input_ids:", input_ids.shape, input_ids.dtype)
print("attention_mask:", attention_mask.shape, attention_mask.dtype)
print("token_type_ids:", token_type_ids.shape, token_type_ids.dtype)

若使用Keras Input层构建输入，需明确指定shape和dtype：

input_ids = tf.keras.layers.Input(shape=(50,), dtype=tf.int32, name='input_ids')
attention_mask = tf.keras.layers.Input(shape=(50,), dtype=tf.int32, name='attention_mask')
token_type_ids = tf.keras.layers.Input(shape=(50,), dtype=tf.int32, name='token_type_ids')

4. 清理缓存并重载预训练模型 #

旧的预训练模型缓存可能和新版本TF不兼容，删除缓存后重新下载：

# 删除Keras缓存的BERT模型
rm -rf ~/.keras/models/bert-*
# 清理pip缓存并重新安装依赖
pip cache purge
pip install --upgrade tensorflow transformers

原因说明 #

该错误源于TensorFlow 2.15.0在图执行模式下对Assert操作的TypeSpec生成逻辑调整，而旧版本transformers库生成的TFBertEmbeddings层包含的断言操作未适配此变化，导致类型不匹配。升级transformers到兼容版本是最根本的解决方式，既能保留GPU支持，又能避免Graph Execution Error。

内容的提问来源于stack exchange，提问作者Tanjim Taharat Aurpa