生成批量分词字符串的ELMo嵌入时遇形状不兼容及Function call stack: pruned错误的排查与解决

问题原因分析 #

你遇到的InvalidArgumentError: Incompatible shapes错误并不是嵌入张量过大导致的，核心问题在于sequence_len参数的形状与ELMo模块内部的计算逻辑不兼容。

在你的代码中，sequence_len被转换成了形状为[4]的一维Tensor，但结合你使用的TensorFlow 2.2 + tf-hub 0.12.0环境，ELMo的tokens签名期望sequence_len的形状能够支持广播计算（比如二维的[4,1]），或者直接传入Python原生列表而非Tensor张量——后者更不容易触发形状匹配问题。

修复后的可运行代码 #

这里提供两种有效的修复方案，推荐用第一种：

import tensorflow as tf
import tensorflow_hub as hub

# 加载ELMo v3模块
elmo = hub.load('https://tfhub.dev/google/elmo/3').signatures['tokens']

# 保持原有的tokens输入结构
tokens = tf.convert_to_tensor(
 [[b'fetal', b'derived', b'definitive', b'erythrocyte', b'', b'', b'', b'', b''],
 [b'splenic', b'red', b'pulp', b'macrophage', b'', b'', b'', b'', b''],
 [b'juxtaglomerular', b'complex', b'cell', b'', b'', b'', b'', b'', b''],
 [b'epithelial', b'cell', b'of', b'large', b'intestine', b'', b'', b'', b'']], tf.string)

# 方案1：直接传入Python列表作为sequence_len（推荐，兼容性更好）
embeddings = elmo(tokens=tokens, sequence_len=[4, 4, 3, 5])['elmo']

# 方案2：将sequence_len扩展为二维Tensor（也能解决问题）
# lens = tf.expand_dims(tf.convert_to_tensor([4, 4, 3, 5], tf.int32), axis=1)
# embeddings = elmo(tokens=tokens, sequence_len=lens)['elmo']

# 验证输出形状，应该得到(4, 9, 1024)
print(embeddings.shape)

避免"Function call stack: pruned"错误的方法 #

这个错误意味着TensorFlow截断了部分堆栈信息，导致你看不到完整的错误溯源，你可以通过以下方法避免或排查：

• 优先用原生Python类型传参：像sequence_len这类简单的参数，直接传列表/元组比转成Tensor更不容易触发形状兼容问题。
• 严格匹配模块输入要求：确认输入张量的形状完全符合ELMo模块的要求——tokens需要是[batch_size, max_sequence_length]的字符串张量，每个样本的长度要统一，空字符串作为padding是符合要求的。
• 缩小调试范围：遇到错误时先尝试用单个样本测试，逐步排查形状问题，避免批量数据干扰。
• 开启详细日志：可以启用TensorFlow的调试日志（tf.debugging.experimental.enable_dump_debug_info），获取完整的错误堆栈，避免被"pruned"截断。

内容的提问来源于stack exchange，提问作者Lorcán