添加Keras自定义层后如何冻结Hugging Face预训练Transformer基础模型

该需求完全可以实现，核心操作是将BERT模型的可训练状态关闭，具体操作如下：

具体步骤 #

1. 加载预训练BERT模型后，新增冻结参数的配置：

bert_model.trainable = False

这行代码会标记BERT内部所有权重不需要在训练过程中更新，仅会计算梯度更新后续你新增的全连接层参数。
2. 完整可运行的代码示例：

import tensorflow as tf
from transformers import TFAutoModel

# 加载预训练BERT，替换为你实际使用的预训练权重版本
bert_model = TFAutoModel.from_pretrained("bert-base-uncased")
# 冻结BERT所有参数
bert_model.trainable = False

# 输入层定义，max_len替换为你实际的序列最大长度
max_len = 128
input_ids = tf.keras.layers.Input(shape=(max_len,), dtype=tf.int32)
attention_masks = tf.keras.layers.Input(shape=(max_len,), dtype=tf.int32)

# BERT输出提取
bert_output = bert_model([input_ids, attention_masks])
pooler_output = bert_output[1] # 提取<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>对应的池化输出
# 新增的回归层
hidden = tf.keras.layers.Dense(100, activation='relu')(pooler_output)
# 回归任务最终输出层
reg_output = tf.keras.layers.Dense(1, activation='linear')(hidden)

# 组装完整模型
model = tf.keras.Model(inputs=[input_ids, attention_masks], outputs=reg_output)
# 编译模型，回归任务一般用MSE损失
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4), loss='mse')

补充说明 #

• 后续如果需要对BERT做微调，只要重新设置bert_model.trainable = True，再重新编译模型即可，你也可以选择只解冻BERT顶部的几层，减少计算量同时避免过拟合。
• 冻结BERT训练时学习率可以比全模型微调稍高，1e-4左右的学习率通常比全微调常用的1e-5收敛效果更好。

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