迁移至TF2后，estimator.export_saved_model为何报变量缺失？

成因分析 #

你遇到的节点名称带数字后缀（如_21）的问题，本质是TF2的层命名机制和分布式训练策略共同作用的结果：

• TF2中，当模型构建过程中出现重复使用同一层实例（比如循环调用同一个Dense层），或者使用tf.distribute.Strategy进行多GPU训练时，TensorFlow会自动给层名称添加数字后缀，避免命名冲突。
• TF1的Estimator在TF2环境下的兼容性问题，也会导致变量命名逻辑和原生TF1不一致，进一步加剧命名不匹配的情况。

解决方法 #

1. 显式指定层名称（推荐） #

在构建学生BERT模型时，给所有关键层（如embedding模块的dense层）手动设置name参数，强制固定节点名称，避免TF自动添加后缀：

# 示例：构建embedding的dense层时指定name
dense_layer = tf.keras.layers.Dense(
    units=hidden_size,
    activation=None,
    name="dense"  # 显式指定名称，避免自动加后缀
)

这样模型节点名称会严格和你定义的一致，和TF1时代的命名逻辑对齐。

2. 修正分布式训练的模型构建逻辑 #

如果使用多GPU训练的tf.distribute.MirroredStrategy，确保在策略范围内只构建一次模型，且所有层都显式指定名称：

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    # 在策略域内一次性构建完整模型，所有关键层指定name
    student_model = build_student_bert(name="bert")  # 自定义模型构建函数，内部层显式命名

避免在策略域重复初始化模型或层，否则TF会为每个副本生成带后缀的层名称。

3. 手动映射变量名加载checkpoint #

如果已经生成了带后缀的checkpoint，无法重新训练生成新的checkpoint，可以通过变量名映射来强制加载：

import re
import tensorflow as tf

# 加载学生模型
student_model = build_student_bert()

# 获取checkpoint中所有变量名
checkpoint_path = "path/to/your/checkpoint"
checkpoint_vars = tf.train.list_variables(checkpoint_path)

# 构建名称映射：带后缀的旧名 -> 模型中的变量
name_map = {}
for var_name, _ in checkpoint_vars:
    # 移除数字后缀（如把dense_21/bias改成dense/bias）
    cleaned_name = re.sub(r'_(\d+)(?=/)', '', var_name)
    # 找到模型中对应的变量
    # 把路径式名称转为属性访问（如bert/embeddings/dense/bias -> bert.embeddings.dense.bias）
    attr_path = cleaned_name.replace('/', '.')
    var = eval(f"student_model.{attr_path}")
    name_map[var_name] = var

# 加载映射后的变量
checkpoint = tf.train.Checkpoint(**name_map)
checkpoint.restore(checkpoint_path).assert_consumed()

注意：eval使用需谨慎，若变量路径复杂，可改用tf.keras.Model.get_layer逐层获取。

4. 替换Estimator导出方式 #

TF2中Estimator属于兼容特性，建议迁移到原生Keras的模型导出方式，避免命名兼容问题：

# 用TF2原生方式导出模型
tf.saved_model.save(student_model, "path/to/exported_model")

后续加载也使用tf.saved_model.load，更适配TF2的生态。

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