TensorFlow 2.x加载ELMo时hub.Module属性错误及适配问题

问题重现 #

我在Jupyter Notebook中尝试用ELMo提取特征，运行代码：

!pip install tensorflow
!pip install tensorflow-hub
import tensorflow_hub as hub
import tensorflow as tf
elmo = hub.Module("https://tfhub.dev/google/elmo/2", trainig = True)

触发错误：

AttributeError: module 'tensorflow_hub' has no attribute 'Module'

环境版本：TensorFlow 2.15.0，TensorFlow Hub 0.16.1，尝试降级TensorFlow无效。

多次尝试后的错误 #

尝试1：改用tf.compat.v1.Module #

运行代码：

import tensorflow as tf
import tensorflow_hub as hub    

elmo = tf.compat.v1.Module("https://tfhub.dev/google/elmo/3", trainable=True)

错误：

TypeError: Module.__init__() got an unexpected keyword argument 'trainable'

尝试2：移除trainable参数 #

运行代码：

elmo = tf.compat.v1.Module("https://tfhub.dev/google/elmo/3")

错误：

ValueError: 'https://tfhub.dev/google/elmo/3' is not a valid module name. Module names must be valid Python identifiers (e.g. a valid class name).

核心原因与解决方案 #

问题出在TensorFlow版本兼容上：hub.Module是TensorFlow 1.x的API，TF2.x版本的TensorFlow Hub已经弃用该接口，改用hub.load()或hub.KerasLayer来加载模型。

方案1：基础特征提取（无需微调） #

import tensorflow as tf
import tensorflow_hub as hub

# 加载预训练ELMo模型
elmo_model = hub.load("https://tfhub.dev/google/elmo/3")

# 定义特征提取函数
def extract_elmo_features(sentences):
    # 传入字符串列表，调用模型的default签名获取特征
    embeddings = elmo_model.signatures["default"](tf.constant(sentences))
    # 可选输出：word_emb（基础词嵌入）、lstm_outputs1（第一层LSTM输出）、lstm_outputs2（第二层LSTM输出）、elmo_output（加权融合后的特征）
    return embeddings["elmo_output"]

# 测试
test_sentences = ["这是一条测试语句", "如何在TF2.x中使用ELMo"]
features = extract_elmo_features(test_sentences)
print("特征形状：", features.shape)

方案2：构建可微调的Keras层 #

如果需要对ELMo进行微调，推荐使用hub.KerasLayer：

import tensorflow as tf
import tensorflow_hub as hub

# 创建可训练的ELMo层
elmo_trainable_layer = hub.KerasLayer(
    handle="https://tfhub.dev/google/elmo/3",
    trainable=True,
    signature="default",
    output_key="elmo_output"
)

# 构建自定义模型
model = tf.keras.Sequential([
    elmo_trainable_layer,
    tf.keras.layers.GlobalAveragePooling1D(),
    tf.keras.layers.Dense(16, activation="relu"),
    tf.keras.layers.Dense(2, activation="softmax")
])

# 测试输入
sample_input = tf.constant(["测试微调模型的输入"])
sample_output = model(sample_input)
print("模型输出形状：", sample_output.shape)

另外注意你最初的代码里有拼写错误：trainig应该是trainable，不过在TF2.x的加载方式中，该参数的设置方式已经改变，上述方案中已经正确处理。

内容的提问来源于stack exchange，提问作者user23483084